zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

add save yaml of opt and hyp to tensorboard log_dir in train()

This commit is contained in:
Alex Stoken 2020-06-16 15:06:13 -05:00
parent a85e6d0fc0
commit d9f446cd81
1 changed files with 10 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

11
train.py
View File

@ -48,7 +48,6 @@ hyp = {'lr0': 0.01, # initial learning rate (SGD=1E-2, Adam=1E-3)
#print(hyp) #print(hyp)
# Overwrite hyp with hyp*.txt (optional) # Overwrite hyp with hyp*.txt (optional)
f = glob.glob('hyp*.txt')
if f: if f:
print('Using %s' % f[0]) print('Using %s' % f[0])
for k, v in zip(hyp.keys(), np.loadtxt(f[0])): for k, v in zip(hyp.keys(), np.loadtxt(f[0])):
@ -64,6 +63,9 @@ def train(hyp):
batch_size = opt.batch_size # 64 batch_size = opt.batch_size # 64
weights = opt.weights # initial training weights weights = opt.weights # initial training weights
#write all results to the tb log_dir, so all data from one run is together
log_dir = tb_writer.log_dir
# Configure # Configure
init_seeds(1) init_seeds(1)
with open(opt.data) as f: with open(opt.data) as f:
@ -192,6 +194,13 @@ def train(hyp):
model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) # attach class weights
model.names = data_dict['names'] model.names = data_dict['names']
#save hyperparamter and training options in run folder
with open(os.path.join(log_dir, 'hyp.yaml', 'w')) as f:
yaml.dump(hyp, f)
with open(os.path.join(log_dir, 'opt.yaml', 'w')) as f:
yaml.dump(opt, f)
# Class frequency # Class frequency
labels = np.concatenate(dataset.labels, 0) labels = np.concatenate(dataset.labels, 0)
c = torch.tensor(labels[:, 0]) # classes c = torch.tensor(labels[:, 0]) # classes