This commit is contained in:
Glenn Jocher 2020-03-11 20:45:14 -07:00
parent e76d4d0ffc
commit 75e88561cb
1 changed files with 2 additions and 2 deletions

View File

@ -240,13 +240,13 @@ def train():
# Hyperparameter Burn-in # Hyperparameter Burn-in
n_burn = 200 # number of burn-in batches n_burn = 200 # number of burn-in batches
if ni <= n_burn: if ni <= n_burn:
g = (ni / n_burn) ** 4 # gain g = (ni / n_burn) ** 2 # gain
for x in model.named_modules(): for x in model.named_modules():
if x[0].endswith('BatchNorm2d'): if x[0].endswith('BatchNorm2d'):
# x[1].momentum = 1 - 0.9 * g # momentum falls from 1 - 0.1 # x[1].momentum = 1 - 0.9 * g # momentum falls from 1 - 0.1
x[1].track_running_stats = ni == n_burn x[1].track_running_stats = ni == n_burn
for x in optimizer.param_groups: for x in optimizer.param_groups:
x['lr'] = x['initial_lr'] * g # gain rises from 0 - 1 x['lr'] = x['initial_lr'] * lf(epoch) * g # gain rises from 0 - 1
# Plot images with bounding boxes # Plot images with bounding boxes
if ni < 1: if ni < 1: