diff --git a/train.py b/train.py index f543da2..be931ea 100644 --- a/train.py +++ b/train.py @@ -52,6 +52,12 @@ def train(hyp): best = wdir + 'best.pt' 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 batch_size = opt.batch_size # 64 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.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 labels = np.concatenate(dataset.labels, 0) c = torch.tensor(labels[:, 0]) # classes