astoken commited on
Commit
de19165
·
1 Parent(s): 2d396be

Fix yaml saving (don't sort keys), reorder --opt keys, bug fix hyp dict accessor

Browse files
Files changed (1) hide show
  1. train.py +7 -6
train.py CHANGED
@@ -91,7 +91,7 @@ def train(hyp):
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,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
- parser.add_argument('--hyp', type=str, default='', help ='hyp cfg path [*.yaml].')
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