hot fix for MultiStepLR_HotFix

kakaobrain · Apr 11, 2020 · 2424224 · 2424224
1 parent 67a2654
commit 2424224
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/FastAutoAugment/lr_scheduler.py b/FastAutoAugment/lr_scheduler.py
@@ -1,5 +1,5 @@
 import torch
-
+from torch.optim.lr_scheduler import MultiStepLR
 from theconf import Config as C
 
 
@@ -10,8 +10,14 @@ def adjust_learning_rate_resnet(optimizer):
     """
 
     if C.get()['epoch'] == 90:
-        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [30, 60, 80])
+        return MultiStepLR_HotFix(optimizer, [30, 60, 80])
     elif C.get()['epoch'] == 270:   # autoaugment
-        return torch.optim.lr_scheduler.MultiStepLR(optimizer, [90, 180, 240])
+        return MultiStepLR_HotFix(optimizer, [90, 180, 240])
     else:
         raise ValueError('invalid epoch=%d for resnet scheduler' % C.get()['epoch'])
+
+
+class MultiStepLR_HotFix(MultiStepLR):
+    def __init__(self, optimizer, milestones, gamma=0.1, last_epoch=-1):
+        super(MultiStepLR_HotFix, self).__init__(optimizer, milestones, gamma, last_epoch)
+        self.milestones = list(milestones)