update to use torch.optim.lr_scheduler.LRScheduler

Yanghan Wang · facebook-github-bot · commit 210809711838 · 2022-12-20T17:13:38.000-08:00
Summary: Pull Request resolved: #4709 pytorch/pytorch#88503 introduces the public version `LRScheduler`, however `isinstance(self.scheduler, torch.optim.lr_scheduler._LRScheduler)` doesn't work anymore because of https://github.com/pytorch/pytorch/blob/1ea11ecb2eea99eb552603b7cf5fbfc59659832d/torch/optim/lr_scheduler.py#L166-L169. It's a bit tricky to make it BC compatible for torch version <= 1.13. V1 of this diff uses try catch block to import the `LRScheduler` and make it available in `detectron2.solver`, then the whole D2 (11528ce) uses this version of `LRScheduler`. There're two drawbacks though: - it adds a little mental burden to figure out what's D2 (11528ce083dc9ff83ee3a8f9086a1ef54d2a402f)'s `LRScheduler`, previously it's clear that the `LRScheduler`/`_LRScheduler` is from `torch`. - it has a name collision with `hooks.LRScheduler`, eg. in the `hooks.py` I have to do `LRScheduler as _LRScheduler`. But I couldn't found a better solution, maybe use try catch block in every file? Reviewed By: sstsai-adl Differential Revision: D42111273 fbshipit-source-id: 0269127de1ba3ef90225c5dfe085bb209f6cf4d1
diff --git a/detectron2/engine/hooks.py b/detectron2/engine/hooks.py
@@ -21,6 +21,7 @@
 import detectron2.utils.comm as comm
 from detectron2.evaluation.testing import flatten_results_dict
 from detectron2.solver import LRMultiplier
+from detectron2.solver import LRScheduler as _LRScheduler
 from detectron2.utils.events import EventStorage, EventWriter
 from detectron2.utils.file_io import PathManager
 
@@ -362,12 +363,12 @@ def scheduler(self):
         return self._scheduler or self.trainer.scheduler
 
     def state_dict(self):
-        if isinstance(self.scheduler, torch.optim.lr_scheduler._LRScheduler):
+        if isinstance(self.scheduler, _LRScheduler):
             return self.scheduler.state_dict()
         return {}
 
     def load_state_dict(self, state_dict):
-        if isinstance(self.scheduler, torch.optim.lr_scheduler._LRScheduler):
+        if isinstance(self.scheduler, _LRScheduler):
             logger = logging.getLogger(__name__)
             logger.info("Loading scheduler from state_dict ...")
             self.scheduler.load_state_dict(state_dict)
diff --git a/detectron2/solver/__init__.py b/detectron2/solver/__init__.py
@@ -1,5 +1,11 @@
 # Copyright (c) Facebook, Inc. and its affiliates.
 from .build import build_lr_scheduler, build_optimizer, get_default_optimizer_params
-from .lr_scheduler import WarmupCosineLR, WarmupMultiStepLR, LRMultiplier, WarmupParamScheduler
+from .lr_scheduler import (
+    LRMultiplier,
+    LRScheduler,
+    WarmupCosineLR,
+    WarmupMultiStepLR,
+    WarmupParamScheduler,
+)
 
 __all__ = [k for k in globals().keys() if not k.startswith("_")]
diff --git a/detectron2/solver/build.py b/detectron2/solver/build.py
@@ -15,7 +15,7 @@
 from detectron2.config import CfgNode
 from detectron2.utils.env import TORCH_VERSION
 
-from .lr_scheduler import LRMultiplier, WarmupParamScheduler
+from .lr_scheduler import LRMultiplier, LRScheduler, WarmupParamScheduler
 
 _GradientClipperInput = Union[torch.Tensor, Iterable[torch.Tensor]]
 _GradientClipper = Callable[[_GradientClipperInput], None]
@@ -267,9 +267,7 @@ def reduce_param_groups(params: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
     return ret
 
 
-def build_lr_scheduler(
-    cfg: CfgNode, optimizer: torch.optim.Optimizer
-) -> torch.optim.lr_scheduler._LRScheduler:
+def build_lr_scheduler(cfg: CfgNode, optimizer: torch.optim.Optimizer) -> LRScheduler:
     """
     Build a LR scheduler from config.
     """
diff --git a/detectron2/solver/lr_scheduler.py b/detectron2/solver/lr_scheduler.py
@@ -11,6 +11,11 @@
     ParamScheduler,
 )
 
+try:
+    from torch.optim.lr_scheduler import LRScheduler
+except ImportError:
+    from torch.optim.lr_scheduler import _LRScheduler as LRScheduler
+
 logger = logging.getLogger(__name__)
 
 
@@ -52,7 +57,7 @@ def __init__(
         )
 
 
-class LRMultiplier(torch.optim.lr_scheduler._LRScheduler):
+class LRMultiplier(LRScheduler):
     """
     A LRScheduler which uses fvcore :class:`ParamScheduler` to multiply the
     learning rate of each param in the optimizer.
@@ -95,7 +100,7 @@ def __init__(
     ):
         """
         Args:
-            optimizer, last_iter: See ``torch.optim.lr_scheduler._LRScheduler``.
+            optimizer, last_iter: See ``torch.optim.lr_scheduler.LRScheduler``.
                 ``last_iter`` is the same as ``last_epoch``.
             multiplier: a fvcore ParamScheduler that defines the multiplier on
                 every LR of the optimizer
@@ -132,7 +137,7 @@ def get_lr(self) -> List[float]:
 # MultiStepLR with WarmupLR but the current LRScheduler design doesn't allow it.
 
 
-class WarmupMultiStepLR(torch.optim.lr_scheduler._LRScheduler):
+class WarmupMultiStepLR(LRScheduler):
     def __init__(
         self,
         optimizer: torch.optim.Optimizer,
@@ -171,7 +176,7 @@ def _compute_values(self) -> List[float]:
         return self.get_lr()
 
 
-class WarmupCosineLR(torch.optim.lr_scheduler._LRScheduler):
+class WarmupCosineLR(LRScheduler):
     def __init__(
         self,
         optimizer: torch.optim.Optimizer,
diff --git a/projects/DeepLab/deeplab/build_solver.py b/projects/DeepLab/deeplab/build_solver.py
@@ -2,14 +2,13 @@
 import torch
 
 from detectron2.config import CfgNode
+from detectron2.solver import LRScheduler
 from detectron2.solver import build_lr_scheduler as build_d2_lr_scheduler
 
 from .lr_scheduler import WarmupPolyLR
 
 
-def build_lr_scheduler(
-    cfg: CfgNode, optimizer: torch.optim.Optimizer
-) -> torch.optim.lr_scheduler._LRScheduler:
+def build_lr_scheduler(cfg: CfgNode, optimizer: torch.optim.Optimizer) -> LRScheduler:
     """
     Build a LR scheduler from config.
     """
diff --git a/projects/DeepLab/deeplab/lr_scheduler.py b/projects/DeepLab/deeplab/lr_scheduler.py
@@ -3,7 +3,7 @@
 from typing import List
 import torch
 
-from detectron2.solver.lr_scheduler import _get_warmup_factor_at_iter
+from detectron2.solver.lr_scheduler import LRScheduler, _get_warmup_factor_at_iter
 
 # NOTE: PyTorch's LR scheduler interface uses names that assume the LR changes
 # only on epoch boundaries. We typically use iteration based schedules instead.
@@ -14,7 +14,7 @@
 # MultiStepLR with WarmupLR but the current LRScheduler design doesn't allow it.
 
 
-class WarmupPolyLR(torch.optim.lr_scheduler._LRScheduler):
+class WarmupPolyLR(LRScheduler):
     """
     Poly learning rate schedule used to train DeepLab.
     Paper: DeepLab: Semantic Image Segmentation with Deep Convolutional Nets,
