Update networks.py

yawudede · Sep 4, 2018 · 65e3da2 · 65e3da2
1 parent bcfbb81
commit 65e3da2
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/models/networks.py b/models/networks.py
@@ -31,6 +31,8 @@ def lambda_rule(epoch):
         scheduler = lr_scheduler.StepLR(optimizer, step_size=opt.lr_decay_iters, gamma=0.1)
     elif opt.lr_policy == 'plateau':
         scheduler = lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.2, threshold=0.01, patience=5)
+    elif opt.lr_policy == 'cosine':
+        scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=opt.niter, eta_min=0)
     else:
         return NotImplementedError('learning rate policy [%s] is not implemented', opt.lr_policy)
     return scheduler