Fixed torch LRScheduler issue and fixed CI (#2780)

* Fixed torch LRScheduler issue and fixed CI

Fixes #2773

* Fixed mypy issues

Author: vfdev-5

.github/workflows/unit-tests.yml

@@ -96,7 +96,9 @@ jobs:
           bash ./tests/run_code_style.sh lint
 
       - name: Run Mypy
-        if: ${{ matrix.os == 'ubuntu-latest'}}
+        # https://github.com/pytorch/ignite/pull/2780
+        # 
+        if: ${{ matrix.os == 'ubuntu-latest' && matrix.pytorch-channel == 'pytorch-nightly'}}
         run: |
           bash ./tests/run_code_style.sh mypy
 

docs/source/conf.py

@@ -342,6 +342,7 @@ def run(self):
     ("py:class", "torch.utils.data.sampler.BatchSampler"),
     ("py:class", "torch.cuda.amp.grad_scaler.GradScaler"),
     ("py:class", "torch.optim.lr_scheduler._LRScheduler"),
+    ("py:class", "torch.optim.lr_scheduler.LRScheduler"),
     ("py:class", "torch.utils.data.dataloader.DataLoader"),
 ]
 

ignite/contrib/engines/common.py

@@ -5,10 +5,15 @@
 
 import torch
 import torch.nn as nn
-from torch.optim.lr_scheduler import _LRScheduler
 from torch.optim.optimizer import Optimizer
 from torch.utils.data.distributed import DistributedSampler
 
+# https://github.com/pytorch/ignite/issues/2773
+try:
+    from torch.optim.lr_scheduler import LRScheduler as PyTorchLRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import _LRScheduler as PyTorchLRScheduler
+
 import ignite.distributed as idist
 from ignite.contrib.handlers import (
     ClearMLLogger,
@@ -37,7 +42,7 @@ def setup_common_training_handlers(
     to_save: Optional[Mapping] = None,
     save_every_iters: int = 1000,
     output_path: Optional[str] = None,
-    lr_scheduler: Optional[Union[ParamScheduler, _LRScheduler]] = None,
+    lr_scheduler: Optional[Union[ParamScheduler, PyTorchLRScheduler]] = None,
     with_gpu_stats: bool = False,
     output_names: Optional[Iterable[str]] = None,
     with_pbars: bool = True,
@@ -140,7 +145,7 @@ def _setup_common_training_handlers(
     to_save: Optional[Mapping] = None,
     save_every_iters: int = 1000,
     output_path: Optional[str] = None,
-    lr_scheduler: Optional[Union[ParamScheduler, _LRScheduler]] = None,
+    lr_scheduler: Optional[Union[ParamScheduler, PyTorchLRScheduler]] = None,
     with_gpu_stats: bool = False,
     output_names: Optional[Iterable[str]] = None,
     with_pbars: bool = True,
@@ -160,9 +165,9 @@ def _setup_common_training_handlers(
         trainer.add_event_handler(Events.ITERATION_COMPLETED, TerminateOnNan())
 
     if lr_scheduler is not None:
-        if isinstance(lr_scheduler, torch.optim.lr_scheduler._LRScheduler):
+        if isinstance(lr_scheduler, PyTorchLRScheduler):
             trainer.add_event_handler(
-                Events.ITERATION_COMPLETED, lambda engine: cast(_LRScheduler, lr_scheduler).step()
+                Events.ITERATION_COMPLETED, lambda engine: cast(PyTorchLRScheduler, lr_scheduler).step()
             )
         else:
             trainer.add_event_handler(Events.ITERATION_STARTED, lr_scheduler)
@@ -226,7 +231,7 @@ def _setup_common_distrib_training_handlers(
     to_save: Optional[Mapping] = None,
     save_every_iters: int = 1000,
     output_path: Optional[str] = None,
-    lr_scheduler: Optional[Union[ParamScheduler, _LRScheduler]] = None,
+    lr_scheduler: Optional[Union[ParamScheduler, PyTorchLRScheduler]] = None,
     with_gpu_stats: bool = False,
     output_names: Optional[Iterable[str]] = None,
     with_pbars: bool = True,

ignite/handlers/param_scheduler.py

@@ -10,9 +10,15 @@
 from typing import Any, cast, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
 
 import torch
-from torch.optim.lr_scheduler import _LRScheduler, ReduceLROnPlateau
+from torch.optim.lr_scheduler import ReduceLROnPlateau
 from torch.optim.optimizer import Optimizer
 
+# https://github.com/pytorch/ignite/issues/2773
+try:
+    from torch.optim.lr_scheduler import LRScheduler as PyTorchLRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import _LRScheduler as PyTorchLRScheduler
+
 from ignite.engine import Engine
 
 
@@ -838,14 +844,15 @@ def print_lr():
 
     def __init__(
         self,
-        lr_scheduler: _LRScheduler,
+        lr_scheduler: PyTorchLRScheduler,
         save_history: bool = False,
         use_legacy: bool = False,
     ):
 
-        if not isinstance(lr_scheduler, _LRScheduler):
+        if not isinstance(lr_scheduler, PyTorchLRScheduler):
             raise TypeError(
-                "Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler._LRScheduler, "
+                "Argument lr_scheduler should be a subclass of "
+                f"torch.optim.lr_scheduler.{PyTorchLRScheduler.__name__}, "
                 f"but given {type(lr_scheduler)}"
             )
 
@@ -882,7 +889,7 @@ def get_param(self) -> Union[float, List[float]]:
 
     @classmethod
     def simulate_values(  # type: ignore[override]
-        cls, num_events: int, lr_scheduler: _LRScheduler, **kwargs: Any
+        cls, num_events: int, lr_scheduler: PyTorchLRScheduler, **kwargs: Any
     ) -> List[List[int]]:
         """Method to simulate scheduled values during num_events events.
 
@@ -894,13 +901,14 @@ def simulate_values(  # type: ignore[override]
             event_index, value
         """
 
-        if not isinstance(lr_scheduler, _LRScheduler):
+        if not isinstance(lr_scheduler, PyTorchLRScheduler):
             raise TypeError(
-                "Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler._LRScheduler, "
+                "Argument lr_scheduler should be a subclass of "
+                f"torch.optim.lr_scheduler.{PyTorchLRScheduler.__name__}, "
                 f"but given {type(lr_scheduler)}"
             )
 
-        # This scheduler uses `torch.optim.lr_scheduler._LRScheduler` which
+        # 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:
@@ -926,7 +934,7 @@ def simulate_values(  # type: ignore[override]
 
 
 def create_lr_scheduler_with_warmup(
-    lr_scheduler: Union[ParamScheduler, _LRScheduler],
+    lr_scheduler: Union[ParamScheduler, PyTorchLRScheduler],
     warmup_start_value: float,
     warmup_duration: int,
     warmup_end_value: Optional[float] = None,
@@ -995,10 +1003,11 @@ def print_lr():
 
     .. versionadded:: 0.4.5
     """
-    if not isinstance(lr_scheduler, (ParamScheduler, _LRScheduler)):
+    if not isinstance(lr_scheduler, (ParamScheduler, PyTorchLRScheduler)):
         raise TypeError(
-            "Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler._LRScheduler or "
-            f"ParamScheduler, but given {type(lr_scheduler)}"
+            "Argument lr_scheduler should be a subclass of "
+            f"torch.optim.lr_scheduler.{PyTorchLRScheduler.__name__} or ParamScheduler, "
+            f"but given {type(lr_scheduler)}"
         )
 
     if not isinstance(warmup_duration, numbers.Integral):
@@ -1018,7 +1027,7 @@ def print_lr():
 
         milestones_values = [(0, warmup_start_value), (warmup_duration - 1, param_group_warmup_end_value)]
 
-        if isinstance(lr_scheduler, _LRScheduler):
+        if isinstance(lr_scheduler, PyTorchLRScheduler):
             init_lr = param_group["lr"]
             if init_lr != param_group_warmup_end_value:
                 milestones_values.append((warmup_duration, init_lr))
@@ -1054,7 +1063,7 @@ def print_lr():
     schedulers = [
         warmup_scheduler,
         lr_scheduler,
-    ]  # type: List[Union[ParamScheduler, ParamGroupScheduler, _LRScheduler]]
+    ]  # type: List[Union[ParamScheduler, ParamGroupScheduler, PyTorchLRScheduler]]
     durations = [milestones_values[-1][0] + 1]
     combined_scheduler = ConcatScheduler(schedulers, durations=durations, save_history=save_history)
 
@@ -1381,7 +1390,9 @@ def load_state_dict(self, state_dict: Mapping) -> None:
             s.load_state_dict(sd)
 
     @classmethod
-    def simulate_values(cls, num_events: int, schedulers: List[_LRScheduler], **kwargs: Any) -> List[List[int]]:
+    def simulate_values(
+        cls, num_events: int, schedulers: List[ParamScheduler], **kwargs: Any
+    ) -> List[List[Union[List[float], float, int]]]:
         """Method to simulate scheduled values during num_events events.
 
         Args:
@@ -1396,7 +1407,7 @@ def simulate_values(cls, num_events: int, schedulers: List[_LRScheduler], **kwar
                 corresponds to the simulated param of scheduler i at 'event_index'th event.
         """
 
-        # This scheduler uses `torch.optim.lr_scheduler._LRScheduler` which
+        # 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:
@@ -1408,9 +1419,9 @@ def simulate_values(cls, num_events: int, schedulers: List[_LRScheduler], **kwar
             torch.save(objs, cache_filepath.as_posix())
 
             values = []
-            scheduler = cls(schedulers=schedulers, **kwargs)  # type: ignore[arg-type]
+            scheduler = cls(schedulers=schedulers, **kwargs)
             for i in range(num_events):
-                params = [scheduler.get_param() for scheduler in schedulers]  # type: ignore[attr-defined]
+                params = [scheduler.get_param() for scheduler in schedulers]
                 values.append([i] + params)
                 scheduler(engine=None)
 

tests/ignite/contrib/engines/test_common.py

@@ -148,8 +148,9 @@ def test_asserts_setup_common_training_handlers():
         train_sampler = MagicMock(spec=DistributedSampler)
         setup_common_training_handlers(trainer, train_sampler=train_sampler)
 
-    with pytest.raises(RuntimeError, match=r"This contrib module requires available GPU"):
-        setup_common_training_handlers(trainer, with_gpu_stats=True)
+    if not torch.cuda.is_available():
+        with pytest.raises(RuntimeError, match=r"This contrib module requires available GPU"):
+            setup_common_training_handlers(trainer, with_gpu_stats=True)
 
     with pytest.raises(TypeError, match=r"Unhandled type of update_function's output."):
         trainer = Engine(lambda e, b: None)

tests/ignite/handlers/test_param_scheduler.py

@@ -621,14 +621,11 @@ def save_lr(engine):
 
 def test_lr_scheduler_asserts():
 
-    with pytest.raises(
-        TypeError, match=r"Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler._LRScheduler"
-    ):
+    err_msg = r"Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler.(_LRScheduler|LRScheduler)"
+    with pytest.raises(TypeError, match=err_msg):
         LRScheduler(123)
 
-    with pytest.raises(
-        TypeError, match=r"Argument lr_scheduler should be a subclass of torch.optim.lr_scheduler._LRScheduler"
-    ):
+    with pytest.raises(TypeError, match=err_msg):
         LRScheduler.simulate_values(1, None)