Repository: pytorch/ignite

Issue:
PiecewiseLinear dict support
In ignite 0.2.1, `PiecewiseLinear` didn't expect the optimizer to be strictly a sub-class of `torch.optim.Optimizer` and allowed generic parameter group [dicts](https://github.com/pytorch/ignite/blob/v0.2.1/ignite/contrib/handlers/param_scheduler.py#L26) as well.

But it looks like in later versions, this generic support was withdrawn.

What's the context behind this change?

I had been successfully using `PiecewiseLinear` from ignite 0.2.1 along with Microsoft's DeepSpeed [ZeRO](https://github.com/microsoft/DeepSpeed/blob/master/deepspeed/pt/deepspeed_zero_optimizer.py#L91) optimizer, which isn't strictly a sub-class of `torch.optim.Optimizer` and Microsoft addressed the reasons for this in their repo. However, due to some useful features in later versions of ignite, we decided to upgrade to a newer version of ignite (0.3 and later also 0.4). In the process of this upgrade, we realized `PiecewiseLinear` no longer works with the ZeRO optimizer.

It would be great if support for generic parameter group dicts is added back.

Discussion:

Patch:
diff --git a/docs/source/contrib/handlers.rst b/docs/source/contrib/handlers.rst
index 188bda71..4560aa9f 100644
--- a/docs/source/contrib/handlers.rst
+++ b/docs/source/contrib/handlers.rst
@@ -75,6 +75,10 @@ mlflow_logger
 tqdm_logger
 -----------
 
+See `tqdm mnist example <https://github.com/pytorch/ignite/blob/master/examples/contrib/mnist/mnist_with_tqdm_logger.py>`_
+for detailed usage.
+
+
 .. automodule:: ignite.contrib.handlers.tqdm_logger
    :members:
 
@@ -88,6 +92,9 @@ polyaxon_logger
 wandb_logger
 ---------------
 
+See `wandb mnist example <https://github.com/pytorch/ignite/blob/master/examples/contrib/mnist/mnist_with_wandb_logger.py>`_
+for detailed usage.
+
 .. automodule:: ignite.contrib.handlers.wandb_logger
    :members:
    :inherited-members:
diff --git a/ignite/contrib/handlers/base_logger.py b/ignite/contrib/handlers/base_logger.py
index ce11b57e..5c6c83f4 100644
--- a/ignite/contrib/handlers/base_logger.py
+++ b/ignite/contrib/handlers/base_logger.py
@@ -1,9 +1,11 @@
 import numbers
 import warnings
 from abc import ABCMeta, abstractmethod
+from collections.abc import Sequence
 from typing import Any, Mapping
 
 import torch
+from torch.optim import Optimizer
 
 from ignite.engine import Engine, State
 
@@ -20,9 +22,13 @@ class BaseOptimizerParamsHandler(BaseHandler):
     """
 
     def __init__(self, optimizer, param_name="lr", tag=None):
-        if not isinstance(optimizer, torch.optim.Optimizer):
+        if not (
+            isinstance(optimizer, Optimizer)
+            or (hasattr(optimizer, "param_groups") and isinstance(optimizer.param_groups, Sequence))
+        ):
             raise TypeError(
-                "Argument optimizer should be of type torch.optim.Optimizer, " "but given {}".format(type(optimizer))
+                "Argument optimizer should be torch.optim.Optimizer or has attribute 'param_groups' as list/tuple, "
+                "but given {}".format(type(optimizer))
             )
 
         self.optimizer = optimizer
diff --git a/ignite/contrib/handlers/mlflow_logger.py b/ignite/contrib/handlers/mlflow_logger.py
index c92d417f..4166111e 100644
--- a/ignite/contrib/handlers/mlflow_logger.py
+++ b/ignite/contrib/handlers/mlflow_logger.py
@@ -167,7 +167,8 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler):
             )
 
     Args:
-        optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): parameter name
         tag (str, optional): common title for all produced plots. For example, 'generator'
     """
diff --git a/ignite/contrib/handlers/neptune_logger.py b/ignite/contrib/handlers/neptune_logger.py
index 43b1fc82..3fc1ca2a 100644
--- a/ignite/contrib/handlers/neptune_logger.py
+++ b/ignite/contrib/handlers/neptune_logger.py
@@ -191,7 +191,8 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler):
             )
 
     Args:
-        optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): parameter name
         tag (str, optional): common title for all produced plots. For example, "generator"
     """
diff --git a/ignite/contrib/handlers/param_scheduler.py b/ignite/contrib/handlers/param_scheduler.py
index 96589bd7..5e1e2434 100644
--- a/ignite/contrib/handlers/param_scheduler.py
+++ b/ignite/contrib/handlers/param_scheduler.py
@@ -18,7 +18,8 @@ class ParamScheduler(metaclass=ABCMeta):
     training.
 
     Args:
-        optimizer (`torch.optim.Optimizer`): optimizer
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): name of optimizer's parameter to update.
         save_history (bool, optional): whether to log the parameter values to
             `engine.state.param_history`, (default=False).
@@ -33,8 +34,14 @@ class ParamScheduler(metaclass=ABCMeta):
 
     def __init__(self, optimizer, param_name, save_history=False, param_group_index=None):
 
-        if not isinstance(optimizer, Optimizer):
-            raise TypeError("Argument optimizer should be torch.optim.Optimizer")
+        if not (
+            isinstance(optimizer, Optimizer)
+            or (hasattr(optimizer, "param_groups") and isinstance(optimizer.param_groups, Sequence))
+        ):
+            raise TypeError(
+                "Argument optimizer should be torch.optim.Optimizer or has attribute 'param_groups' as list/tuple, "
+                "but given {}".format(type(optimizer))
+            )
 
         self.optimizer = optimizer
         self.param_group_index = param_group_index
@@ -212,7 +219,8 @@ class CyclicalScheduler(ParamScheduler):
     cycle of some size.
 
     Args:
-        optimizer (`torch.optim.Optimizer`): optimizer
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): name of optimizer's parameter to update.
         start_value (float): value at start of cycle.
         end_value (float): value at the middle of the cycle.
@@ -287,7 +295,8 @@ class LinearCyclicalScheduler(CyclicalScheduler):
     adjusts it back to 'start_value' for a half-cycle.
 
     Args:
-        optimizer (`torch.optim.Optimizer`): optimizer
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): name of optimizer's parameter to update.
         start_value (float): value at start of cycle.
         end_value (float): value at the middle of the cycle.
@@ -332,7 +341,8 @@ class CosineAnnealingScheduler(CyclicalScheduler):
     wave (as suggested in [Smith17]_).
 
     Args:
-        optimizer (`torch.optim.Optimizer`): optimizer
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): name of optimizer's parameter to update.
         start_value (float): value at start of cycle.
         end_value (float): value at the end of the cycle.
@@ -818,7 +828,8 @@ class PiecewiseLinear(ParamScheduler):
     Piecewise linear parameter scheduler
 
     Args:
-        optimizer (`torch.optim.Optimizer`): optimizer.
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): name of optimizer's parameter to update.
         milestones_values (list of tuples (int, float)): list of tuples (event index, parameter value)
             represents milestones and parameter. Milestones should be increasing integers.
diff --git a/ignite/contrib/handlers/polyaxon_logger.py b/ignite/contrib/handlers/polyaxon_logger.py
index 55c8410b..e675cbee 100644
--- a/ignite/contrib/handlers/polyaxon_logger.py
+++ b/ignite/contrib/handlers/polyaxon_logger.py
@@ -154,7 +154,8 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler):
             )
 
     Args:
-        optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): parameter name
         tag (str, optional): common title for all produced plots. For example, "generator"
     """
diff --git a/ignite/contrib/handlers/tensorboard_logger.py b/ignite/contrib/handlers/tensorboard_logger.py
index 0ded2dfb..99d61041 100644
--- a/ignite/contrib/handlers/tensorboard_logger.py
+++ b/ignite/contrib/handlers/tensorboard_logger.py
@@ -163,7 +163,8 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler):
             )
 
     Args:
-        optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): parameter name
         tag (str, optional): common title for all produced plots. For example, "generator"
     """
diff --git a/ignite/contrib/handlers/trains_logger.py b/ignite/contrib/handlers/trains_logger.py
index b0b35e25..a1fd811a 100644
--- a/ignite/contrib/handlers/trains_logger.py
+++ b/ignite/contrib/handlers/trains_logger.py
@@ -182,7 +182,8 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler):
             )
 
     Args:
-        optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): parameter name
         tag (str, optional): common title for all produced plots. For example, "generator"
     """
diff --git a/ignite/contrib/handlers/visdom_logger.py b/ignite/contrib/handlers/visdom_logger.py
index f8002e35..a4b9a337 100644
--- a/ignite/contrib/handlers/visdom_logger.py
+++ b/ignite/contrib/handlers/visdom_logger.py
@@ -218,7 +218,8 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler, _BaseVisDrawer):
             )
 
     Args:
-        optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): parameter name
         tag (str, optional): common title for all produced plots. For example, "generator"
         show_legend (bool, optional): flag to show legend in the window
diff --git a/ignite/contrib/handlers/wandb_logger.py b/ignite/contrib/handlers/wandb_logger.py
index 8123d8f3..dd494578 100644
--- a/ignite/contrib/handlers/wandb_logger.py
+++ b/ignite/contrib/handlers/wandb_logger.py
@@ -166,7 +166,8 @@ class OptimizerParamsHandler(BaseOptimizerParamsHandler):
             )
 
     Args:
-        optimizer (torch.optim.Optimizer): torch optimizer which parameters to log
+        optimizer (torch.optim.Optimizer or object): torch optimizer or any object with attribute ``param_groups``
+            as a sequence.
         param_name (str): parameter name
         tag (str, optional): common title for all produced plots. For example, "generator"
         sync (bool, optional): If set to False, process calls to log in a seperate thread. Default (None) uses whatever


Test patch:
diff --git a/tests/ignite/contrib/handlers/__init__.py b/tests/ignite/contrib/handlers/__init__.py
index e69de29b..f9cd4aeb 100644
--- a/tests/ignite/contrib/handlers/__init__.py
+++ b/tests/ignite/contrib/handlers/__init__.py
@@ -0,0 +1,14 @@
+class MockFP16DeepSpeedZeroOptimizer:
+    def __init__(self, optimizer):
+        self.optimizer = optimizer
+
+    def step(self, closure=None):
+        self.optimizer.step()
+
+    def _get_param_groups(self):
+        return self.optimizer.param_groups
+
+    def _set_param_groups(self, value):
+        self.optimizer.param_groups = value
+
+    param_groups = property(_get_param_groups, _set_param_groups)
diff --git a/tests/ignite/contrib/handlers/test_base_logger.py b/tests/ignite/contrib/handlers/test_base_logger.py
index 315b6de0..4ca342de 100644
--- a/tests/ignite/contrib/handlers/test_base_logger.py
+++ b/tests/ignite/contrib/handlers/test_base_logger.py
@@ -7,6 +7,7 @@ import torch
 from ignite.contrib.handlers import CustomPeriodicEvent
 from ignite.contrib.handlers.base_logger import BaseLogger, BaseOptimizerParamsHandler, BaseOutputHandler
 from ignite.engine import Engine, Events, State
+from tests.ignite.contrib.handlers import MockFP16DeepSpeedZeroOptimizer
 
 
 class DummyOutputHandler(BaseOutputHandler):
@@ -15,8 +16,13 @@ class DummyOutputHandler(BaseOutputHandler):
 
 
 class DummyOptParamsHandler(BaseOptimizerParamsHandler):
-    def __call__(self, *args, **kwargs):
-        pass
+    def __call__(self, engine, logger, event_name, **kwargs):
+        tag_prefix = "{}/".format(self.tag) if self.tag else ""
+        params = {
+            "{}{}/group_{}".format(tag_prefix, self.param_name, i): float(param_group[self.param_name])
+            for i, param_group in enumerate(self.optimizer.param_groups)
+        }
+        return params
 
 
 class DummyLogger(BaseLogger):
@@ -41,6 +47,9 @@ def test_base_output_handler_wrong_setup():
     with pytest.raises(TypeError, match="global_step_transform should be a function"):
         DummyOutputHandler("tag", metric_names=["loss"], global_step_transform="abc")
 
+    with pytest.raises(TypeError, match=r"Argument optimizer should be torch.optim.Optimizer"):
+        DummyOptParamsHandler({}, "lr")
+
 
 def test_base_output_handler_setup_output_metrics():
 
@@ -81,6 +90,16 @@ def test_base_output_handler_setup_output_metrics():
     assert metrics == true_metrics
 
 
+def test_opt_params_handler_on_non_torch_optimizers():
+    tensor = torch.zeros([1], requires_grad=True)
+    base_optimizer = torch.optim.SGD([tensor], lr=0.1234)
+    optimizer = MockFP16DeepSpeedZeroOptimizer(base_optimizer)
+    handler = DummyOptParamsHandler(optimizer=optimizer, param_name="lr")
+    res = handler(engine=None, logger=None, event_name=None)
+    assert isinstance(res, dict)
+    assert "lr/group_0" in res and res["lr/group_0"] == 0.1234
+
+
 def test_attach():
 
     n_epochs = 5
diff --git a/tests/ignite/contrib/handlers/test_param_scheduler.py b/tests/ignite/contrib/handlers/test_param_scheduler.py
index c8521af0..8e32aa59 100644
--- a/tests/ignite/contrib/handlers/test_param_scheduler.py
+++ b/tests/ignite/contrib/handlers/test_param_scheduler.py
@@ -1,3 +1,5 @@
+from unittest.mock import MagicMock
+
 import numpy as np
 import pytest
 import torch
@@ -14,6 +16,7 @@ from ignite.contrib.handlers.param_scheduler import (
     create_lr_scheduler_with_warmup,
 )
 from ignite.engine import Engine, Events
+from tests.ignite.contrib.handlers import MockFP16DeepSpeedZeroOptimizer
 
 try:
     from torch.optim.lr_scheduler import MultiplicativeLR
@@ -26,10 +29,12 @@ else:
     has_multiplicative_lr = LooseVersion(torch.__version__) >= LooseVersion("1.5.0")
 
 
+class FakeParamScheduler(ParamScheduler):
+    def get_param(self):
+        return [0]
+
+
 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)
@@ -46,6 +51,9 @@ def test_param_scheduler_asserts():
     with pytest.raises(ValueError, match=r"Required state attribute 'event_index' is absent in provided state_dict"):
         lr_scheduler.load_state_dict({})
 
+    with pytest.raises(TypeError, match=r"Argument optimizer should be torch.optim.Optimizer"):
+        FakeParamScheduler({}, "lr")
+
 
 def test_linear_scheduler():
 
@@ -1220,3 +1228,32 @@ def test_scheduler_with_param_groups():
 
     torch_lr_scheduler = StepLR(optimizer, step_size=50, gamma=0.5)
     _test(LRScheduler(torch_lr_scheduler), optimizer)
+
+
+def test_lr_scheduling_on_non_torch_optimizers():
+    # tests https://github.com/pytorch/ignite/issues/1162
+    optimizer = MagicMock()
+    optimizer.param_groups = [{"params": 0}]
+    FakeParamScheduler(optimizer, "lr")
+
+    tensor = torch.zeros([1], requires_grad=True)
+    base_optimizer = torch.optim.SGD([tensor], lr=0)
+    optimizer = MockFP16DeepSpeedZeroOptimizer(base_optimizer)
+
+    milestones_values = [(5, 0.5), (15, 1.0)]
+
+    scheduler = PiecewiseLinear(optimizer, "lr", milestones_values=milestones_values)
+
+    def save_lr(engine):
+        lrs.append(optimizer.param_groups[0]["lr"])
+
+    trainer = Engine(lambda engine, batch: None)
+    trainer.add_event_handler(Events.ITERATION_COMPLETED, scheduler)
+    trainer.add_event_handler(Events.ITERATION_COMPLETED, save_lr)
+
+    lrs = []
+    trainer.run([0] * 15, max_epochs=1)
+
+    assert lrs == list(
+        map(pytest.approx, [0.5, 0.5, 0.5, 0.5, 0.5, 0.5, 0.55, 0.6, 0.65, 0.7, 0.75, 0.8, 0.85, 0.9, 0.95,],)
+    )