add ut case

open-mmlab · Oct 28, 2022 · 152c029 · 152c029
1 parent eebadc3
commit 152c029
Showing 1 changed file with 18 additions and 0 deletions.
diff --git a/tests/test_optim/test_scheduler/test_param_scheduler.py b/tests/test_optim/test_scheduler/test_param_scheduler.py
@@ -438,6 +438,24 @@ def test_cosine_restart_scheduler(self):
             eta_min=0)
         self._test_scheduler_value(scheduler, targets, epochs=10)
 
+        epochs = 10
+        t = 10
+        eta_min = 1e-7
+        single_targets = [
+            eta_min + (0.05 - eta_min) * (1 + math.cos(math.pi * x / t)) / 2
+            for x in range(epochs)
+        ]
+        targets = [
+            single_targets, [x * self.layer2_mult for x in single_targets]
+        ]
+        scheduler = CosineRestartParamScheduler(
+            self.optimizer,
+            param_name='lr',
+            periods=[t],
+            restart_weights=[1],
+            eta_min=eta_min)
+        self._test_scheduler_value(scheduler, targets, epochs=10)
+
     def _check_scheduler_state_dict(self, construct, construct2, epochs=10):
         scheduler = construct()
         for _ in range(epochs):