Update schedulers.py

jiaowoguanren0615 · Jan 1, 2024 · 36cc9b8 · 36cc9b8
1 parent a1fa8a7
commit 36cc9b8
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/SegFormer/utils/schedulers.py b/SegFormer/utils/schedulers.py
@@ -1,11 +1,11 @@
 import torch
 import math
-from torch.optim.lr_scheduler import LRScheduler, LambdaLR
+from torch.optim.lr_scheduler import _LRScheduler, LambdaLR
 
 __all__ = ['PolyLR', 'WarmupPolyLR', 'WarmupCosineLR', 'WarmupExpLR', 'WarmupLR']
 
 
-class PolyLR(LRScheduler):
+class PolyLR(_LRScheduler):
     def __init__(self, optimizer, max_iter, decay_iter=1, power=0.9, last_epoch=-1) -> None:
         self.decay_iter = decay_iter
         self.max_iter = max_iter
@@ -20,7 +20,7 @@ def get_lr(self):
             return [factor * lr for lr in self.base_lrs]
 
 
-class WarmupLR(LRScheduler):
+class WarmupLR(_LRScheduler):
     def __init__(self, optimizer, warmup_iter=500, warmup_ratio=5e-4, warmup='exp', last_epoch=-1) -> None:
         self.warmup_iter = warmup_iter
         self.warmup_ratio = warmup_ratio
@@ -144,4 +144,4 @@ def f(x):
 
     plt.plot(np.arange(len(lrs)), np.array(lrs))
     plt.grid()
-    plt.show()
+    plt.show()