add warmup

main_moco.py

@@ -390,8 +390,11 @@ class ProgressMeter(object):
 
 
 def adjust_learning_rate(optimizer, init_lr, epoch, args):
-    """Decays the learning rate with half-cycle cosine"""
-    lr = init_lr * 0.5 * (1. + math.cos(math.pi * epoch / args.epochs))
+    """Decays the learning rate with half-cycle cosine after warmup"""
+    if epoch < args.warmup_epochs:
+        lr = init_lr / args.warmup_epochs * epoch
+    else:
+        lr = init_lr * 0.5 * (1. + math.cos(math.pi * (epoch - warmup_epochs) / (args.epochs - args.warmup_epochs)))
     for param_group in optimizer.param_groups:
         param_group['lr'] = lr
 

moco/optimizer.py

@@ -11,9 +11,9 @@ class LARS(torch.optim.Optimizer):
     """
     LARS optimizer, no rate scaling or weight decay for normalization layers and biases.
     """
-    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, trust_coefficient=0.001):
+    def __init__(self, params, lr=0, weight_decay=0, momentum=0.9, nesterov=False, trust_coefficient=0.001):
         defaults = dict(lr=lr, weight_decay=weight_decay, 
-                        momentum=momentum, nesterov=False,
+                        momentum=momentum, nesterov=nesterov,
                         trust_coefficient=trust_coefficient)
         super().__init__(params, defaults)