Repository: pytorch/ignite

Issue:
CosineAnnealingWarmRestarts is not compatible with LRScheduler.simulate_values()
## 🐛 Bug description

`CosineAnnealingWarmRestarts` is not compatible with `LRScheduler.simulate_values()`. Precisely, a object of type `CosineAnnealingWarmRestarts` can't be replicated by `LRScheduler._replicate_lr_scheduler()`. It works for `CosineAnnealingLR`.

For example : 

```
lr_scheduler = CosineAnnealingWarmRestarts(optimizer=optimizer, T_0=10, eta_min=1e-3)
lr_values = LRScheduler.simulate_values(num_events=50, lr_scheduler=lr_scheduler)
```

produces the following error 

```
Traceback (most recent call last):
  File "tutorials/misc/lr_schedulers.py", line 56, in <module>
    lr_values = LRScheduler.simulate_values(num_events=50, lr_scheduler=lr_scheduler)
  File "/work/desrozis/Softwares/conda/envs/focus-light/lib/python3.7/site-packages/ignite/contrib/handlers/param_scheduler.py", line 606, in simulate_values
    copy_lr_scheduler = LRScheduler._replicate_lr_scheduler(lr_scheduler)
  File "/work/desrozis/Softwares/conda/envs/focus-light/lib/python3.7/site-packages/ignite/contrib/handlers/param_scheduler.py", line 627, in _replicate_lr_scheduler
    copy_lr_scheduler = lr_scheduler_cls(optimizer=dummy_optimizer, **kwargs)
TypeError: __init__() got an unexpected keyword argument 'T_i'
```

This issue is due to the assumption that `lr_scheduler.__state_dict__` contains arguments of method `CosineAnnealingWarmRestarts.__init__()`. A workaround could be to remove `T_i` in a similar way to `base_lrs` and `last_epoch` but it's not very satisfactory...

Hope it helps to have a better and stonger ignite.

## Environment

 - PyTorch Version (e.g., 1.4): 1.4
 - Ignite Version (e.g., 0.3.0): 0.3.0
 - OS (e.g., Linux): Linux
 - How you installed Ignite (`conda`, `pip`, source): conda
 - Python version: 3.7.6
 - Any other relevant information:



Discussion:

Patch:
diff --git a/ignite/contrib/handlers/param_scheduler.py b/ignite/contrib/handlers/param_scheduler.py
index 57e460b4..96589bd7 100644
--- a/ignite/contrib/handlers/param_scheduler.py
+++ b/ignite/contrib/handlers/param_scheduler.py
@@ -1,9 +1,11 @@
 import math
 import numbers
+import tempfile
 from abc import ABCMeta, abstractmethod
 from collections import OrderedDict
 from collections.abc import Mapping, Sequence
 from copy import copy
+from pathlib import Path
 from typing import List, Optional, Union
 
 import torch
@@ -449,7 +451,16 @@ class ConcatScheduler(ParamScheduler):
 
         self.schedulers = schedulers
         self.durations = durations
-        super(ConcatScheduler, self).__init__(optimizer=_get_fake_optimizer(), param_name="", save_history=save_history)
+
+        self.optimizer = self.schedulers[0].optimizer
+        if not (all(id(s.optimizer) == id(self.optimizer) for s in self.schedulers)):
+            raise ValueError("schedulers should be related to same optimizer")
+
+        # schedulers should have save_history sync with ParamGroupScheduler
+        for s in schedulers:
+            s.save_history = save_history
+
+        super(ConcatScheduler, self).__init__(optimizer=self.optimizer, param_name="", save_history=save_history)
 
         self._scheduler_index = 0
         self._current_scheduler = None
@@ -510,7 +521,6 @@ class ConcatScheduler(ParamScheduler):
         if self._current_duration == 0:
             self._scheduler_index += 1
             self._setup_scheduler()
-
         self._current_scheduler(engine, name)
         self._current_duration -= 1
 
@@ -549,22 +559,40 @@ class ConcatScheduler(ParamScheduler):
         """
         if param_names is not None and not isinstance(param_names, (list, tuple)):
             raise ValueError("Argument param_names should be list or tuple of strings")
-        output = []
-
-        # Need to copy all schedulers otherwise unsafe
-        copy_schedulers = [_replicate_scheduler(s) for s in schedulers]
 
-        scheduler = cls(copy_schedulers, durations, save_history=False)
-        if param_names is None:
-            param_names = [scheduler.param_name]
-        for i in range(num_events):
-            scheduler(engine=None)
-            values = [i]
-            for param_name in param_names:
-                params = [p[param_name] for p in scheduler.optimizer_param_groups]
-                values = values + params
-            output.append(values)
-        return output
+        # This scheduler uses `ParamScheduler` which
+        # should be replicated in order to simulate LR values and
+        # not perturb original scheduler.
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            cache_filepath = Path(tmpdirname) / "ignite_lr_scheduler_cache.pt"
+            objs = {"lr_scheduler_{}".format(i): s.state_dict() for i, s in enumerate(schedulers)}
+            # all schedulers should be related to the same optimizer
+            objs["optimizer"] = schedulers[0].optimizer.state_dict()
+
+            torch.save(objs, cache_filepath.as_posix())
+
+            # do not save_history
+            for s in schedulers:
+                s.save_history = False
+
+            output = []
+            scheduler = cls(schedulers=schedulers, save_history=False, durations=durations, **kwargs)
+            if param_names is None:
+                param_names = [scheduler.param_name]
+            for i in range(num_events):
+                scheduler(engine=None)
+                values = [i]
+                for param_name in param_names:
+                    params = [p[param_name] for p in scheduler.optimizer_param_groups]
+                    values = values + params
+                output.append(values)
+
+            objs = torch.load(cache_filepath.as_posix())
+            for i, s in enumerate(schedulers):
+                s.load_state_dict(objs["lr_scheduler_{}".format(i)])
+                s.optimizer.load_state_dict(objs["optimizer"])
+
+            return output
 
 
 class LRScheduler(ParamScheduler):
@@ -591,7 +619,7 @@ class LRScheduler(ParamScheduler):
         trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
     """
 
-    def __init__(self, lr_scheduler, save_history=False, **kwds):
+    def __init__(self, lr_scheduler, save_history=False, **kwargs):
 
         if not isinstance(lr_scheduler, _LRScheduler):
             raise TypeError(
@@ -635,33 +663,36 @@ class LRScheduler(ParamScheduler):
             list of pairs: [event_index, value]
 
         """
+
+        if not isinstance(lr_scheduler, _LRScheduler):
+            raise TypeError(
+                "Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler._LRScheduler, "
+                "but given {}".format(type(lr_scheduler))
+            )
+
         # This scheduler uses `torch.optim.lr_scheduler._LRScheduler` which
         # should be replicated in order to simulate LR values and
         # not perturb original scheduler.
-        copy_lr_scheduler = LRScheduler._replicate_lr_scheduler(lr_scheduler)
-        values = []
-        scheduler = cls(save_history=False, lr_scheduler=copy_lr_scheduler)
-        for i in range(num_events):
-            params = [p[scheduler.param_name] for p in scheduler.optimizer_param_groups]
-            values.append([i] + params)
-            scheduler(engine=None)
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            cache_filepath = Path(tmpdirname) / "ignite_lr_scheduler_cache.pt"
+            obj = {
+                "lr_scheduler": lr_scheduler.state_dict(),
+                "optimizer": lr_scheduler.optimizer.state_dict(),
+            }
+            torch.save(obj, cache_filepath.as_posix())
 
-        return values
+            values = []
+            scheduler = cls(save_history=False, lr_scheduler=lr_scheduler, **kwargs)
+            for i in range(num_events):
+                params = [p[scheduler.param_name] for p in scheduler.optimizer_param_groups]
+                values.append([i] + params)
+                scheduler(engine=None)
 
-    @staticmethod
-    def _replicate_lr_scheduler(lr_scheduler):
-        if not isinstance(lr_scheduler, _LRScheduler):
-            raise TypeError("lr_scheduler should inherit of _LRScheduler, got {}".format(type(lr_scheduler)))
-        lr_scheduler_cls = lr_scheduler.__class__
-        dummy_optimizer = _replicate_optimizer(lr_scheduler.optimizer)
-        for group in dummy_optimizer.param_groups:
-            group.setdefault("initial_lr", group["lr"])
-        kwargs = lr_scheduler.state_dict()
-        for k in [_k for _k in kwargs.keys() if "_" == _k[0]] + ["base_lrs", "last_epoch"]:
-            del kwargs[k]
-        copy_lr_scheduler = lr_scheduler_cls(optimizer=dummy_optimizer, **kwargs)
-        copy_lr_scheduler.load_state_dict(lr_scheduler.state_dict())
-        return copy_lr_scheduler
+            obj = torch.load(cache_filepath.as_posix())
+            lr_scheduler.load_state_dict(obj["lr_scheduler"])
+            lr_scheduler.optimizer.load_state_dict(obj["optimizer"])
+
+            return values
 
 
 def create_lr_scheduler_with_warmup(
@@ -741,7 +772,7 @@ def create_lr_scheduler_with_warmup(
             if init_lr != param_group_warmup_end_value:
                 milestones_values.append((warmup_duration, init_lr))
 
-            lr_scheduler = LRScheduler(lr_scheduler)
+            lr_scheduler = LRScheduler(lr_scheduler, save_history=save_history)
         else:
             init_lr = lr_scheduler.get_param()
             if init_lr == param_group_warmup_end_value:
@@ -913,16 +944,29 @@ class ParamGroupScheduler(ParamScheduler):
         self.schedulers = schedulers
         self.names = names
 
-        optimizer = self.schedulers[0].optimizer
-        if not (all(id(s.optimizer) == id(optimizer) for s in schedulers)):
+        self.optimizer = self.schedulers[0].optimizer
+        if not (all(id(s.optimizer) == id(self.optimizer) for s in schedulers)):
             raise ValueError("schedulers should be related to same optimizer")
 
-        super(ParamGroupScheduler, self).__init__(optimizer=optimizer, param_name="lr", save_history=save_history)
+        # schedulers should have save_history sync with ParamGroupScheduler
+        for s in schedulers:
+            s.save_history = save_history
+
+        super(ParamGroupScheduler, self).__init__(optimizer=self.optimizer, param_name="lr", save_history=save_history)
 
     def __call__(self, engine, name=None):
         for scheduler, name in zip(self.schedulers, self.names):
             scheduler(engine, name)
 
+    @property
+    def save_history(self):
+        return self.schedulers[0].save_history
+
+    @save_history.setter
+    def save_history(self, value):
+        for s in self.schedulers:
+            s.save_history = value
+
     def get_param(self) -> Union[List[float], float]:
         return [scheduler.get_param() for scheduler in self.schedulers]
 
@@ -981,35 +1025,31 @@ class ParamGroupScheduler(ParamScheduler):
             list of pairs: [event_index, value]
 
         """
-        copy_lr_schedulers = [_replicate_scheduler(s) for s in schedulers]
-        values = []
-        scheduler = cls(schedulers=copy_lr_schedulers)
-        for i in range(num_events):
-            scheduler(engine=None)
-            params = scheduler.get_param()
-            values.append([i] + params)
-        return values
 
+        # This scheduler uses `torch.optim.lr_scheduler._LRScheduler` which
+        # should be replicated in order to simulate LR values and
+        # not perturb original scheduler.
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            cache_filepath = Path(tmpdirname) / "ignite_lr_scheduler_cache.pt"
+            objs = {"lr_scheduler_{}".format(i): s.state_dict() for i, s in enumerate(schedulers)}
+            # all schedulers should be related to the same optimizer
+            objs["optimizer"] = schedulers[0].optimizer.state_dict()
+
+            torch.save(objs, cache_filepath.as_posix())
 
-def _replicate_scheduler(scheduler, save_history=False):
-    if isinstance(scheduler, LRScheduler):
-        return LRScheduler(LRScheduler._replicate_lr_scheduler(scheduler.lr_scheduler), save_history=save_history)
-    elif isinstance(scheduler, ConcatScheduler):
-        copy_schedulers = [_replicate_scheduler(s, save_history=save_history) for s in scheduler.schedulers]
-        return ConcatScheduler(copy_schedulers, durations=scheduler.durations, save_history=save_history)
-    elif isinstance(scheduler, ParamGroupScheduler):
-        copy_optimizer = _replicate_optimizer(scheduler.optimizer)
-        copy_schedulers = [_replicate_scheduler(s, save_history=save_history) for s in scheduler.schedulers]
-        for s in copy_schedulers:
-            s.optimizer = copy_optimizer
-        return ParamGroupScheduler(schedulers=copy_schedulers, names=scheduler.names, save_history=save_history)
-    elif isinstance(scheduler, ParamScheduler):
-        new_scheduler = copy(scheduler)
-        new_scheduler.optimizer = _replicate_optimizer(new_scheduler.optimizer)
-        new_scheduler.save_history = save_history
-        return new_scheduler
-    else:
-        raise TypeError("Unknown scheduler type {}".format(type(scheduler)))
+            values = []
+            scheduler = cls(schedulers=schedulers, **kwargs)
+            for i in range(num_events):
+                params = scheduler.get_param()
+                values.append([i] + params)
+                scheduler(engine=None)
+
+            objs = torch.load(cache_filepath.as_posix())
+            for i, s in enumerate(schedulers):
+                s.load_state_dict(objs["lr_scheduler_{}".format(i)])
+                s.optimizer.load_state_dict(objs["optimizer"])
+
+            return values
 
 
 def _get_fake_optimizer(optimizer_cls=None, **kwargs):
@@ -1018,19 +1058,3 @@ def _get_fake_optimizer(optimizer_cls=None, **kwargs):
         optimizer_cls = torch.optim.SGD
         kwargs["lr"] = 0.01
     return optimizer_cls([t], **kwargs)
-
-
-def _replicate_optimizer(optimizer):
-    cls = optimizer.__class__
-    defaults = copy(optimizer.defaults)
-    if not isinstance(defaults["lr"], numbers.Real):
-        defaults["lr"] = 0.01
-    param_groups = optimizer.param_groups
-    param_groups = [
-        # do no copy params
-        {k: v for k, v in pg.items() if k != "params"}
-        for pg in param_groups
-    ]
-    for pg in param_groups:
-        pg["params"] = torch.zeros([1], requires_grad=True)
-    return cls(param_groups, **defaults)


Test patch:
diff --git a/tests/ignite/contrib/handlers/test_param_scheduler.py b/tests/ignite/contrib/handlers/test_param_scheduler.py
index d8c965d6..7ea0af0a 100644
--- a/tests/ignite/contrib/handlers/test_param_scheduler.py
+++ b/tests/ignite/contrib/handlers/test_param_scheduler.py
@@ -1,7 +1,7 @@
 import numpy as np
 import pytest
 import torch
-from torch.optim.lr_scheduler import ExponentialLR, StepLR
+from torch.optim.lr_scheduler import ExponentialLR, MultiplicativeLR, StepLR
 
 from ignite.contrib.handlers.param_scheduler import (
     ConcatScheduler,
@@ -9,12 +9,34 @@ from ignite.contrib.handlers.param_scheduler import (
     LinearCyclicalScheduler,
     LRScheduler,
     ParamGroupScheduler,
+    ParamScheduler,
     PiecewiseLinear,
     create_lr_scheduler_with_warmup,
 )
 from ignite.engine import Engine, Events
 
 
+def test_param_scheduler_asserts():
+    class FakeParamScheduler(ParamScheduler):
+        def get_param(self):
+            return [0]
+
+    t1 = torch.zeros([1], requires_grad=True)
+    t2 = torch.zeros([1], requires_grad=True)
+    optimizer = torch.optim.SGD([{"params": t1, "lr": 0.1}, {"params": t2, "lr": 0.1}])
+
+    lr_scheduler = FakeParamScheduler(optimizer, "lr")
+
+    with pytest.raises(RuntimeError, match=r"size of value is different than optimizer_param_groups"):
+        lr_scheduler(None)
+
+    with pytest.raises(TypeError, match=r"Argument state_dict should be a dictionary, but given"):
+        lr_scheduler.load_state_dict(None)
+
+    with pytest.raises(ValueError, match=r"Required state attribute 'event_index' is absent in provided state_dict"):
+        lr_scheduler.load_state_dict({})
+
+
 def test_linear_scheduler():
 
     with pytest.raises(TypeError, match=r"Argument optimizer should be torch.optim.Optimizer"):
@@ -285,6 +307,12 @@ def test_concat_scheduler_asserts():
             num_events=123, schedulers=[scheduler_1, scheduler_2], durations=[15], param_names="abc"
         )
 
+    optimizer_2 = torch.optim.SGD([tensor], lr=0)
+    scheduler_3 = CosineAnnealingScheduler(optimizer_2, "lr", start_value=0.0, end_value=1.0, cycle_size=10)
+
+    with pytest.raises(ValueError, match=r"schedulers should be related to same optimizer"):
+        ConcatScheduler([scheduler_1, scheduler_3], durations=[30,])
+
 
 def test_concat_scheduler_state_dict():
     tensor = torch.zeros([1], requires_grad=True)
@@ -313,6 +341,9 @@ def test_concat_scheduler_state_dict():
     with pytest.raises(ValueError, match=r"Input state_dict contains 0 state_dicts of concatenated schedulers"):
         concat_scheduler.load_state_dict({"schedulers": []})
 
+    with pytest.raises(TypeError, match=r"Argument state_dict should be a dictionary, but given"):
+        concat_scheduler.load_state_dict(None)
+
 
 def test_concat_scheduler_two_schedulers():
     tensor = torch.zeros([1], requires_grad=True)
@@ -563,6 +594,9 @@ def test_lr_scheduler_asserts():
     with pytest.raises(TypeError):
         LRScheduler(123)
 
+    with pytest.raises(TypeError):
+        LRScheduler.simulate_values(1, None)
+
 
 def test_lr_scheduler():
     def _test(torch_lr_scheduler_cls, **kwargs):
@@ -622,23 +656,11 @@ def test_lr_scheduler():
         )
         assert lrs == pytest.approx([v for i, v in simulated_values])
 
-    _test(torch.optim.lr_scheduler.StepLR, step_size=5, gamma=0.5)
-    _test(torch.optim.lr_scheduler.ExponentialLR, gamma=0.78)
-
-    # test _replicate_lr_scheduler
-    tensor = torch.zeros([1], requires_grad=True)
-    optimizer = torch.optim.SGD([tensor], lr=0.01)
-    lr_scheduler = torch.optim.lr_scheduler.ExponentialLR(optimizer=optimizer, gamma=0.78)
-    init_lr_scheduler_state = dict(lr_scheduler.state_dict())
-    copy_lr_scheduler = LRScheduler._replicate_lr_scheduler(lr_scheduler)
-    for _ in range(10):
-        optimizer.step()
-        lr_scheduler.step()
-
-    assert copy_lr_scheduler.state_dict() == init_lr_scheduler_state
-
-    with pytest.raises(TypeError):
-        LRScheduler._replicate_lr_scheduler(12)
+    _test(StepLR, step_size=5, gamma=0.5)
+    _test(ExponentialLR, gamma=0.78)
+    _test(MultiplicativeLR, lr_lambda=lambda epoch: 0.95)
+    # bug #813
+    # _test(torch.optim.lr_scheduler.CosineAnnealingWarmRestarts, T_0=10, T_mult=10)
 
 
 def test_piecewiselinear_asserts():
@@ -658,6 +680,9 @@ def test_piecewiselinear_asserts():
     with pytest.raises(ValueError):
         PiecewiseLinear(optimizer, "lr", milestones_values=[(10, 0.5), (5, 0.6)])
 
+    with pytest.raises(ValueError):
+        PiecewiseLinear(optimizer, "lr", milestones_values=[(0.5, 1)])
+
 
 def test_piecewiselinear():
     def _test(milestones_as_np_int):
@@ -1085,11 +1110,21 @@ def test_param_group_scheduler_asserts():
     with pytest.raises(ValueError, match=r"Input state_dict contains 0 state_dicts of param group schedulers"):
         scheduler.load_state_dict({"schedulers": []})
 
+    with pytest.raises(ValueError, match=r"Required state attribute 'schedulers' is absent in provided state_dict"):
+        scheduler.load_state_dict({})
+
     with pytest.raises(
         ValueError, match=r"Name of scheduler from input state dict does not " r"correspond to required one"
     ):
         scheduler.load_state_dict({"schedulers": [("a", lr_scheduler1.state_dict()), ("bad_name", {})]})
 
+    optimizer2 = torch.optim.SGD([{"params": t1, "lr": 0.1}, {"params": t2, "lr": 0.1}])
+    lr_scheduler3 = LinearCyclicalScheduler(
+        optimizer2, "lr", param_group_index=0, start_value=1.0, end_value=0.0, cycle_size=10
+    )
+    with pytest.raises(ValueError, match=r"schedulers should be related to same optimizer"):
+        ParamGroupScheduler(schedulers=[lr_scheduler1, lr_scheduler3])
+
 
 def test_param_group_scheduler():
     def _test(lr_schedulers, optimizer):
@@ -1115,6 +1150,10 @@ def test_param_group_scheduler():
             assert [lr[0] for lr in lrs] == pytest.approx([lr[1] for lr in lrs])
             scheduler.load_state_dict(state_dict)
 
+            values = ParamGroupScheduler.simulate_values(max_epochs * num_iterations, lr_schedulers)
+            assert [lr[1] for lr in values] == pytest.approx([lr[2] for lr in values])
+            assert [lr[0] for lr in lrs] == pytest.approx([lr[1] for lr in values])
+
     t1 = torch.zeros([1], requires_grad=True)
     t2 = torch.zeros([1], requires_grad=True)
     optimizer = torch.optim.SGD([{"params": t1, "lr": 0.1}, {"params": t2, "lr": 0.1}])