diff --git a/train.py b/train.py index deb74928..7d394829 100644 --- a/train.py +++ b/train.py @@ -134,11 +134,12 @@ def train(): if opt.transfer or opt.prebias: # transfer learning edge (yolo) layers nf = int(model.module_defs[model.yolo_layers[0] - 1]['filters']) # yolo layer size (i.e. 255) - for p in optimizer.param_groups: - # lower param count allows more aggressive training settings: i.e. SGD ~0.1 lr0, ~0.9 momentum - p['lr'] *= 100 # lr gain - if p.get('momentum') is not None: # for SGD but not Adam - p['momentum'] *= 0.9 + if opt.prebias: + for p in optimizer.param_groups: + # lower param count allows more aggressive training settings: i.e. SGD ~0.1 lr0, ~0.9 momentum + p['lr'] *= 100 # lr gain + if p.get('momentum') is not None: # for SGD but not Adam + p['momentum'] *= 0.9 for p in model.parameters(): if opt.prebias and p.numel() == nf: # train (yolo biases)