Fix yaml saving (don't sort keys), reorder --opt keys, bug fix hyp dict accessor
Browse files
train.py
CHANGED
@@ -91,7 +91,7 @@ def train(hyp):
|
|
91 |
else:
|
92 |
pg0.append(v) # all else
|
93 |
|
94 |
-
if hyp
|
95 |
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) #use default beta2, adjust beta1 for Adam momentum per momentum adjustments in https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
|
96 |
else:
|
97 |
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
|
@@ -190,10 +190,10 @@ def train(hyp):
|
|
190 |
|
191 |
#save hyperparamter and training options in run folder
|
192 |
with open(os.path.join(log_dir, 'hyp.yaml'), 'w') as f:
|
193 |
-
yaml.dump(hyp, f)
|
194 |
|
195 |
with open(os.path.join(log_dir, 'opt.yaml'), 'w') as f:
|
196 |
-
yaml.dump(vars(opt), f)
|
197 |
|
198 |
# Class frequency
|
199 |
labels = np.concatenate(dataset.labels, 0)
|
@@ -370,10 +370,11 @@ def train(hyp):
|
|
370 |
if __name__ == '__main__':
|
371 |
check_git_status()
|
372 |
parser = argparse.ArgumentParser()
|
373 |
-
parser.add_argument('--epochs', type=int, default=300)
|
374 |
-
parser.add_argument('--batch-size', type=int, default=16)
|
375 |
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model cfg path[*.yaml]')
|
376 |
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data cfg path [*.yaml]')
|
|
|
|
|
|
|
377 |
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes. Assumes square imgs.')
|
378 |
parser.add_argument('--rect', action='store_true', help='rectangular training')
|
379 |
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
@@ -386,7 +387,7 @@ if __name__ == '__main__':
|
|
386 |
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
387 |
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
|
388 |
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
|
389 |
-
|
390 |
opt = parser.parse_args()
|
391 |
|
392 |
opt.cfg = check_file(opt.cfg) # check file
|
|
|
91 |
else:
|
92 |
pg0.append(v) # all else
|
93 |
|
94 |
+
if hyp['optimizer'] =='adam':
|
95 |
optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) #use default beta2, adjust beta1 for Adam momentum per momentum adjustments in https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR
|
96 |
else:
|
97 |
optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
|
|
|
190 |
|
191 |
#save hyperparamter and training options in run folder
|
192 |
with open(os.path.join(log_dir, 'hyp.yaml'), 'w') as f:
|
193 |
+
yaml.dump(hyp, f, sort_keys=False)
|
194 |
|
195 |
with open(os.path.join(log_dir, 'opt.yaml'), 'w') as f:
|
196 |
+
yaml.dump(vars(opt), f, sort_keys=False)
|
197 |
|
198 |
# Class frequency
|
199 |
labels = np.concatenate(dataset.labels, 0)
|
|
|
370 |
if __name__ == '__main__':
|
371 |
check_git_status()
|
372 |
parser = argparse.ArgumentParser()
|
|
|
|
|
373 |
parser.add_argument('--cfg', type=str, default='models/yolov5s.yaml', help='model cfg path[*.yaml]')
|
374 |
parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data cfg path [*.yaml]')
|
375 |
+
parser.add_argument('--hyp', type=str, default='',help='hyp cfg path [*.yaml].')
|
376 |
+
parser.add_argument('--epochs', type=int, default=300)
|
377 |
+
parser.add_argument('--batch-size', type=int, default=16)
|
378 |
parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes. Assumes square imgs.')
|
379 |
parser.add_argument('--rect', action='store_true', help='rectangular training')
|
380 |
parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
|
|
|
387 |
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
|
388 |
parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%')
|
389 |
parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
|
390 |
+
|
391 |
opt = parser.parse_args()
|
392 |
|
393 |
opt.cfg = check_file(opt.cfg) # check file
|