astoken commited on
Commit
7abf202
1 Parent(s): 7edbf65

Mode all optimizer settings to 'hyp.yaml', integrate proper momentum with Adam optimizer

Browse files
Files changed (1) hide show
  1. train.py +8 -5
train.py CHANGED
@@ -20,8 +20,9 @@ except:
20
 
21
 
22
  # Hyperparameters
23
- hyp = {'lr0': 0.01, # initial learning rate (SGD=1E-2, Adam=1E-3)
24
- 'momentum': 0.937, # SGD momentum
 
25
  'weight_decay': 5e-4, # optimizer weight decay
26
  'giou': 0.05, # giou loss gain
27
  'cls': 0.58, # cls loss gain
@@ -90,8 +91,11 @@ def train(hyp):
90
  else:
91
  pg0.append(v) # all else
92
 
93
- optimizer = optim.Adam(pg0, lr=hyp['lr0']) if opt.adam else \
94
- optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True)
 
 
 
95
  optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
96
  optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
97
  print('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
@@ -380,7 +384,6 @@ if __name__ == '__main__':
380
  parser.add_argument('--weights', type=str, default='', help='initial weights path')
381
  parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')
382
  parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
383
- parser.add_argument('--adam', action='store_true', help='use adam optimizer')
384
  parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%')
385
  parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset')
386
  parser.add_argument('--hyp', type=str, default='', help ='hyp cfg path [*.yaml].')
 
20
 
21
 
22
  # Hyperparameters
23
+ hyp = {'optimizer': 'adam' #if none, default is SGD
24
+ 'lr0': 0.01, # initial learning rate (SGD=1E-2, Adam=1E-3)
25
+ 'momentum': 0.937, # SGD momentum/Adam beta1
26
  'weight_decay': 5e-4, # optimizer weight decay
27
  'giou': 0.05, # giou loss gain
28
  'cls': 0.58, # cls loss gain
 
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)
98
+
99
  optimizer.add_param_group({'params': pg1, 'weight_decay': hyp['weight_decay']}) # add pg1 with weight_decay
100
  optimizer.add_param_group({'params': pg2}) # add pg2 (biases)
101
  print('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0)))
 
384
  parser.add_argument('--weights', type=str, default='', help='initial weights path')
385
  parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied')
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].')