This commit is contained in:
Glenn Jocher 2020-03-11 15:57:37 -07:00
parent 4089735c5e
commit e40d4c87f2
1 changed files with 9 additions and 5 deletions

View File

@ -237,11 +237,15 @@ def train():
imgs = imgs.to(device).float() / 255.0 # uint8 to float32, 0 - 255 to 0.0 - 1.0
targets = targets.to(device)
# Hyperparameter burn-in
# n_burn = 100 # number of burn-in batches
# if ni < n_burn:
# for x in optimizer.param_groups:
# x['lr'] = x['initial_lr'] * (ni / n_burn) ** 4 # gain rises from 0 - 1
# Hyperparameter Burn-in
n_burn = 100 # number of burn-in batches
if ni <= n_burn:
g = (ni / n_burn) ** 4 # gain
for x in model.named_modules():
if x[0].endswith('BatchNorm2d'):
x[1].momentum = 1 - 0.9 * g # momentum falls from 1 - 0.1
for x in optimizer.param_groups:
x['lr'] = x['initial_lr'] * g # gain rises from 0 - 1
# Plot images with bounding boxes
if ni < 1: