diff --git a/train.py b/train.py index a045241f..c455d868 100644 --- a/train.py +++ b/train.py @@ -235,9 +235,10 @@ def train(cfg, # SGD burn-in if epoch == 0 and i <= n_burnin: - lr = hyp['lr0'] * (i / n_burnin) ** 4 + g = (i / n_burnin) ** 4 # gain for x in optimizer.param_groups: - x['lr'] = lr + x['lr'] = hyp['lr0'] * g + x['weight_decay'] = hyp['weight_decay'] * g # Run model pred = model(imgs)