Fixed parameter scheduler bug with CosineAnnealingWarmRestarts

ignite/handlers/param_scheduler.py

@@ -10,7 +10,7 @@
 from typing import Any, cast, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
 
 import torch
-from torch.optim.lr_scheduler import ReduceLROnPlateau
+from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, ReduceLROnPlateau
 from torch.optim.optimizer import Optimizer
 
 # https://github.com/pytorch/ignite/issues/2773
@@ -792,6 +792,64 @@ def simulate_values(  # type: ignore[override]
             return output
 
 
+class _CosineAnnealingWarmRestarts:
+    def __init__(self, lr_scheduler: CosineAnnealingWarmRestarts):
+        self._lr_scheduler = lr_scheduler
+
+    @property
+    def last_epoch(self) -> int:
+        return self._lr_scheduler.last_epoch
+
+    @last_epoch.setter
+    def last_epoch(self, value: int) -> None:
+        self._lr_scheduler.last_epoch = value
+
+    @property
+    def optimizer(self) -> torch.optim.Optimizer:
+        return self._lr_scheduler.optimizer
+
+    def get_lr(self, epoch: Optional[int] = None) -> List[float]:
+        if epoch is None and self.last_epoch < 0:
+            epoch = 0
+        self._lr_scheduler.T_mult: int
+        if epoch is None:
+            epoch = self.last_epoch + 1
+            self._lr_scheduler.T_cur = self._lr_scheduler.T_cur + 1
+            if self._lr_scheduler.T_cur >= self._lr_scheduler.T_i:
+                self._lr_scheduler.T_cur = self._lr_scheduler.T_cur - self._lr_scheduler.T_i
+                self._lr_scheduler.T_i = self._lr_scheduler.T_i * self._lr_scheduler.T_mult
+        else:
+            if epoch < 0:
+                raise ValueError("Expected non-negative epoch, but got {}".format(epoch))
+            if epoch >= self._lr_scheduler.T_0:
+                if self._lr_scheduler.T_mult == 1:
+                    self._lr_scheduler.T_cur = epoch % self._lr_scheduler.T_0
+                else:
+                    n = int(
+                        math.log(
+                            (epoch / self._lr_scheduler.T_0 * (self._lr_scheduler.T_mult - 1) + 1),
+                            self._lr_scheduler.T_mult,
+                        )
+                    )
+                    self._lr_scheduler.T_cur = epoch - self._lr_scheduler.T_0 * (self._lr_scheduler.T_mult**n - 1) / (
+                        self._lr_scheduler.T_mult - 1
+                    )
+                    self._lr_scheduler.T_i = self._lr_scheduler.T_0 * self._lr_scheduler.T_mult ** (n)
+            else:
+                self._lr_scheduler.T_i = self._lr_scheduler.T_0
+                self._lr_scheduler.T_cur = epoch
+
+        self.last_epoch = math.floor(epoch)
+
+        return [
+            self._lr_scheduler.eta_min
+            + (base_lr - self._lr_scheduler.eta_min)
+            * (1 + math.cos(math.pi * self._lr_scheduler.T_cur / self._lr_scheduler.T_i))
+            / 2
+            for base_lr in self._lr_scheduler.base_lrs
+        ]
+
+
 class LRScheduler(ParamScheduler):
     """A wrapper class to call `torch.optim.lr_scheduler` objects as `ignite` handlers.
 
@@ -852,8 +910,12 @@ def __init__(
                 f"torch.optim.lr_scheduler.{PyTorchLRScheduler.__name__}, "
                 f"but given {type(lr_scheduler)}"
             )
+        self.lr_scheduler: Union[PyTorchLRScheduler, _CosineAnnealingWarmRestarts]
+        if isinstance(lr_scheduler, torch.optim.lr_scheduler.CosineAnnealingWarmRestarts):
+            self.lr_scheduler = _CosineAnnealingWarmRestarts(lr_scheduler)
+        else:
+            self.lr_scheduler = lr_scheduler
 
-        self.lr_scheduler = lr_scheduler
         super(LRScheduler, self).__init__(
             optimizer=self.lr_scheduler.optimizer,
             param_name="lr",

tests/ignite/handlers/test_param_scheduler.py

@@ -3,7 +3,7 @@
 import numpy as np
 import pytest
 import torch
-from torch.optim.lr_scheduler import ExponentialLR, StepLR
+from torch.optim.lr_scheduler import CosineAnnealingWarmRestarts, ExponentialLR, StepLR
 
 from ignite.engine import Engine, Events
 from ignite.handlers.param_scheduler import (
@@ -1362,3 +1362,44 @@ def test_reduce_lr_on_plateau_scheduler_asserts():
     with pytest.raises(ValueError, match=r"Length of argument metric_values should be equal to num_events."):
         metric_values = [0.1, 0.2, 0.3, 0.4, 0.5, 0.6]
         ReduceLROnPlateauScheduler.simulate_values(5, metric_values, 0.01)
+
+
+@pytest.mark.parametrize("warmup_end_value", [0.23, None])
+def test_create_lr_scheduler_with_warmup_cosine(warmup_end_value):
+    lr = 0.2
+    steps = 100
+    warm_steps = 50
+    warm_start = 0.023
+
+    def get_optim():
+        t1 = torch.zeros([1], requires_grad=True)
+        return torch.optim.SGD([t1], lr=lr)
+
+    def get_cos_shed():
+        return CosineAnnealingWarmRestarts(optimizer, T_0=12, T_mult=3, verbose=False)
+
+    optimizer = get_optim()
+    scheduler = get_cos_shed()
+    cosine_lrs = []
+    for i in range(steps):
+        cosine_lrs.append(optimizer.param_groups[0]["lr"])
+        scheduler.step()
+
+    optimizer = get_optim()
+    scheduler = create_lr_scheduler_with_warmup(
+        get_cos_shed(), warmup_start_value=warm_start, warmup_end_value=warmup_end_value, warmup_duration=warm_steps
+    )
+
+    warm_lrs = []
+    for epoch in range(warm_steps + steps):
+        scheduler(None)
+        warm_lrs.append(optimizer.param_groups[0]["lr"])
+
+    if warmup_end_value is not None:
+        assert (
+            np.linspace(warm_start, warmup_end_value, warm_steps).round(3) == np.array(warm_lrs[:warm_steps]).round(3)
+        ).all()
+        assert warm_lrs[warm_steps:] == cosine_lrs
+    else:
+        assert (np.linspace(warm_start, lr, warm_steps).round(3) == np.array(warm_lrs[:warm_steps]).round(3)).all()
+        assert warm_lrs[warm_steps - 1 : -1] == cosine_lrs