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

Merge pull request #338 from alexstoken/hyp_save_bugfix 

Move hyp and opt yaml save to top of train()
This commit is contained in:
Glenn Jocher 2020-07-09 14:35:32 -07:00 committed by GitHub
parent bf6f41567a 9d631408a2
commit 2b6209a9d5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 6 additions and 6 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
train.py
View File

@ -52,6 +52,12 @@ def train(hyp):
best = wdir + 'best.pt' best = wdir + 'best.pt'
results_file = log_dir + os.sep + 'results.txt' results_file = log_dir + os.sep + 'results.txt'
# Save run settings
with open(Path(log_dir) / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(Path(log_dir) / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
epochs = opt.epochs # 300 epochs = opt.epochs # 300
batch_size = opt.batch_size # 64 batch_size = opt.batch_size # 64
weights = opt.weights # initial training weights weights = opt.weights # initial training weights
@ -171,12 +177,6 @@ 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 run settings
with open(Path(log_dir) / 'hyp.yaml', 'w') as f:
yaml.dump(hyp, f, sort_keys=False)
with open(Path(log_dir) / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)
# 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