astoken commited on
Commit
52bac22
1 Parent(s): 8b6dbb7

Add in --resume functionality with option to specify path or to get most recent run

Browse files
Files changed (1) hide show
  1. train.py +8 -2
train.py CHANGED
@@ -22,7 +22,7 @@ except:
22
 
23
 
24
  # Hyperparameters
25
- hyp = {'optimizer': 'SGD', # ['adam, 'SGD', None] if none, default is SGD
26
  'lr0': 0.01, # initial learning rate (SGD=1E-2, Adam=1E-3)
27
  'momentum': 0.937, # SGD momentum/Adam beta1
28
  'weight_decay': 5e-4, # optimizer weight decay
@@ -375,7 +375,7 @@ if __name__ == '__main__':
375
  parser.add_argument('--batch-size', type=int, default=16)
376
  parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes. Assumes square imgs.')
377
  parser.add_argument('--rect', action='store_true', help='rectangular training')
378
- parser.add_argument('--resume', action='store_true', help='resume training from last.pt')
379
  parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
380
  parser.add_argument('--notest', action='store_true', help='only test final epoch')
381
  parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
@@ -390,7 +390,13 @@ if __name__ == '__main__':
390
 
391
  opt = parser.parse_args()
392
 
 
 
 
 
393
  opt.weights = last if opt.resume and not opt.weights else opt.weights
 
 
394
  opt.cfg = check_file(opt.cfg) # check file
395
  opt.data = check_file(opt.data) # check file
396
  opt.hyp = check_file(opt.hyp) if opt.hyp else '' #check file
 
22
 
23
 
24
  # Hyperparameters
25
+ hyp = {'optimizer': 'SGD', # ['adam', 'SGD', None] if none, default is SGD
26
  'lr0': 0.01, # initial learning rate (SGD=1E-2, Adam=1E-3)
27
  'momentum': 0.937, # SGD momentum/Adam beta1
28
  'weight_decay': 5e-4, # optimizer weight decay
 
375
  parser.add_argument('--batch-size', type=int, default=16)
376
  parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='train,test sizes. Assumes square imgs.')
377
  parser.add_argument('--rect', action='store_true', help='rectangular training')
378
+ parser.add_argument('--resume', nargs='?', const = 'get_last', default=False, help='resume training from given path/to/last.pt, or most recent run if blank.')
379
  parser.add_argument('--nosave', action='store_true', help='only save final checkpoint')
380
  parser.add_argument('--notest', action='store_true', help='only test final epoch')
381
  parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check')
 
390
 
391
  opt = parser.parse_args()
392
 
393
+ # use given path/to/last.pt or find most recent run if no path given
394
+ last = get_latest_run() if opt.resume == 'get_last' else opt.resume
395
+ if last and not opt.weights:
396
+ print(f'Resuming training from {last}')
397
  opt.weights = last if opt.resume and not opt.weights else opt.weights
398
+
399
+
400
  opt.cfg = check_file(opt.cfg) # check file
401
  opt.data = check_file(opt.data) # check file
402
  opt.hyp = check_file(opt.hyp) if opt.hyp else '' #check file