Make lr scheduler configurable

utils/generic_utils.py

@@ -142,6 +142,20 @@ def lr_decay(init_lr, global_step, warmup_steps):
     return lr
 
 
+class AnnealLR(torch.optim.lr_scheduler._LRScheduler):
+    def __init__(self, optimizer, warmup_steps=0.1):
+        self.warmup_steps = float(warmup_steps)
+        super(AnnealLR, self).__init__(optimizer, last_epoch)
+
+    def get_lr(self):
+        return [
+            base_lr * self.warmup_steps**0.5 * torch.min([
+                self.last_epoch * self.warmup_steps**-1.5, self.last_epoch**
+                -0.5
+            ]) for base_lr in self.base_lrs
+        ]
+
+
 def mk_decay(init_mk, max_epoch, n_epoch):
     return init_mk * ((max_epoch - n_epoch) / max_epoch)