yjh0410 преди 1 година
родител
ревизия
8e5d83e32e
променени са 1 файла, в които са добавени 3 реда и са изтрити 2 реда
  1. 3 2
      odlab/utils/lr_scheduler.py

+ 3 - 2
odlab/utils/lr_scheduler.py

@@ -11,8 +11,9 @@ class LinearWarmUpScheduler(object):
 
     def set_lr(self, optimizer, lr):
         for param_group in optimizer.param_groups:
-            param_group['lr'] = lr
-
+            init_lr = param_group['initial_lr']
+            ratio = init_lr / self.base_lr
+            param_group['lr'] = lr * ratio
 
     def __call__(self, iter, optimizer):
         # warmup