Add LightningModule.lr_scheduler_step

.gitignore

@@ -142,7 +142,7 @@ mnist/
 legacy/checkpoints/
 *.gz
 *ubyte
-
+MNIST/
 
 # pl tests
 ml-runs/

CHANGELOG.md

@@ -58,6 +58,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added a `PrecisionPlugin.teardown` method ([#10990](https://github.com/PyTorchLightning/pytorch-lightning/issues/10990))
 
 
+- Added `LightningModule.lr_scheduler_step` ([#10249](https://github.com/PyTorchLightning/pytorch-lightning/pull/10249))
+
 
 ### Changed
 

docs/source/common/optimizers.rst

@@ -252,6 +252,29 @@ If you want to call schedulers that require a metric value after each epoch, con
 
 -----
 
+Bring your own Custom Learning Rate Schedulers
+----------------------------------------------
+Lightning allows using custom learning rate schedulers that aren't available in `PyTorch natively <https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate>`_.
+One good example is `Timm Schedulers <https://github.com/rwightman/pytorch-image-models/blob/master/timm/scheduler/scheduler.py>`_. When using custom learning rate schedulers
+relying on a different API from Native PyTorch ones, you should override the :meth:`~pytorch_lightning.core.lightning.LightningModule.lr_scheduler_step` with your desired logic.
+If you are using native PyTorch schedulers, there is no need to override this hook since Lightning will handle it optimally by default.
+
+.. code-block:: python
+
+    from timm.scheduler import TanhLRScheduler
+
+
+    def configure_optimizers(self):
+        optimizer = ...
+        scheduler = TanhLRScheduler(optimizer, ...)
+        return [optimizer], [{"scheduler": scheduler, "interval": "epoch"}]
+
+
+    def lr_scheduler_step(self, scheduler, optimizer_idx, metric=None):
+        scheduler.step(epoch=self.current_epoch)  # timm's scheduler need the epoch value
+
+-----
+
 Use closure for LBFGS-like optimizers
 -------------------------------------
 It is a good practice to provide the optimizer with a closure function that performs a ``forward``, ``zero_grad`` and

pytorch_lightning/core/lightning.py

@@ -53,7 +53,7 @@
 from pytorch_lightning.utilities.model_summary import ModelSummary, summarize
 from pytorch_lightning.utilities.parsing import collect_init_args
 from pytorch_lightning.utilities.signature_utils import is_param_in_hook_signature
-from pytorch_lightning.utilities.types import _METRIC_COLLECTION, EPOCH_OUTPUT, STEP_OUTPUT
+from pytorch_lightning.utilities.types import _METRIC_COLLECTION, EPOCH_OUTPUT, LRSchedulerTypeUnion, STEP_OUTPUT
 from pytorch_lightning.utilities.warnings import WarningCache
 
 warning_cache = WarningCache()
@@ -1493,6 +1493,42 @@ def configure_gradient_clipping(self, optimizer, optimizer_idx, gradient_clip_va
             optimizer, gradient_clip_val=gradient_clip_val, gradient_clip_algorithm=gradient_clip_algorithm
         )
 
+    def lr_scheduler_step(
+        self,
+        scheduler: LRSchedulerTypeUnion,
+        optimizer_idx: Optional[int] = None,
+        metric: Optional[Union[float, torch.Tensor]] = None,
+    ) -> None:
+        r"""
+        Override this method to adjust the default way the
+        :class:`~pytorch_lightning.trainer.trainer.Trainer` calls each scheduler.
+        By default, Lightning calls ``step()`` and as shown in the example
+        for each scheduler based on its ``interval``.
+
+        Args:
+            scheduler: Learning rate scheduler.
+            optimizer_idx: Index of the optimizer associated with this scheduler.
+            metric: Value of the metric used for schedulers like ``ReduceLROnPlateau``.
+
+        Examples::
+
+            # DEFAULT
+            def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+                if metric is None:
+                    scheduler.step()
+                else:
+                    scheduler.step(metric)
+
+            # Alternative way to update schedulers if it requires an epoch value
+            def lr_scheduler_step(self, scheduler, optimizer_idx, metric):
+                scheduler.step(epoch=self.current_epoch)
+
+        """
+        if metric is None:
+            scheduler.step()
+        else:
+            scheduler.step(metric)
+
     def optimizer_step(
         self,
         epoch: int,

pytorch_lightning/core/optimizer.py

@@ -19,6 +19,7 @@
 import torch
 from torch import optim
 from torch.optim import Optimizer
+from typing_extensions import Protocol, runtime_checkable
 
 import pytorch_lightning as pl
 from pytorch_lightning.utilities import AMPType, rank_zero_warn
@@ -168,7 +169,9 @@ def closure_dis():
             trainer.strategy.optimizer_step(self._optimizer, self._optimizer_idx, closure, **kwargs)
 
 
-def _init_optimizers_and_lr_schedulers(model: "pl.LightningModule") -> Tuple[List, List, List]:
+def _init_optimizers_and_lr_schedulers(
+    model: "pl.LightningModule",
+) -> Tuple[List[Optimizer], List[Dict[str, Any]], List[int]]:
     """Calls `LightningModule.configure_optimizers` and parses and validates the output."""
     model.trainer._lightning_optimizers = None
     optim_conf = model.trainer._call_lightning_module_hook("configure_optimizers", pl_module=model)
@@ -298,10 +301,13 @@ def _configure_schedulers_automatic_opt(schedulers: list, monitor: Optional[str]
             lr_schedulers.append(
                 {**default_config, "scheduler": scheduler, "reduce_on_plateau": True, "monitor": monitor}
             )
-        elif isinstance(scheduler, optim.lr_scheduler._LRScheduler):
-            lr_schedulers.append({**default_config, "scheduler": scheduler})
         else:
-            raise ValueError(f'The provided lr scheduler "{scheduler}" is invalid')
+            lr_schedulers.append({**default_config, "scheduler": scheduler})
+
+        current_scheduler = lr_schedulers[-1]["scheduler"]
+        if not isinstance(current_scheduler, _SupportedLRScheduler):
+            raise ValueError(f"The provided lr scheduler `{current_scheduler.__class__.__name__}` is invalid.")
+
     return lr_schedulers
 
 
@@ -325,6 +331,11 @@ def _configure_schedulers_manual_opt(schedulers: list, monitor: Optional[str]) -
             lr_schedulers.append({**default_config, **scheduler})
         else:
             lr_schedulers.append({**default_config, "scheduler": scheduler})
+
+        current_scheduler = lr_schedulers[-1]["scheduler"]
+        if not isinstance(current_scheduler, _SupportedLRScheduler):
+            raise ValueError(f"The provided lr scheduler `{current_scheduler.__class__.__name__}` is invalid.")
+
     return lr_schedulers
 
 
@@ -341,7 +352,7 @@ def _get_default_scheduler_config() -> Dict[str, Any]:
     }
 
 
-def _validate_scheduler_optimizer(optimizers: List[Any], lr_schedulers: List[Any]) -> None:
+def _validate_scheduler_optimizer(optimizers: List[Optimizer], lr_schedulers: List[Dict[str, Any]]) -> None:
     if any(sch["scheduler"].optimizer not in optimizers for sch in lr_schedulers):
         raise MisconfigurationException(
             "Some schedulers are attached with an optimizer that wasn't returned from `configure_optimizers`."
@@ -394,3 +405,17 @@ def zero_grad(self, set_to_none: Optional[bool] = False) -> None:
 
     def __repr__(self) -> str:
         return "No Optimizer"
+
+
+@runtime_checkable
+class _SupportedLRScheduler(Protocol):
+    """This class is used to detect if an object is stateful using `isinstance(obj, _SupportedLRScheduler)`"""
+
+    def step(self, *args: Any, **kwargs: Any) -> None:
+        ...
+
+    def state_dict(self) -> Dict[str, Any]:
+        ...
+
+    def load_state_dict(self, state_dict: Dict[str, Any]) -> None:
+        ...

pytorch_lightning/loops/epoch/training_epoch_loop.py

@@ -503,11 +503,12 @@ def _update_learning_rates(
                 self.scheduler_progress.increment_ready()
 
                 # update LR
-                if lr_scheduler["reduce_on_plateau"]:
-                    lr_scheduler["scheduler"].step(monitor_val)
-                else:
-                    lr_scheduler["scheduler"].step()
-
+                self.trainer._call_lightning_module_hook(
+                    "lr_scheduler_step",
+                    lr_scheduler["scheduler"],
+                    optimizer_idx=lr_scheduler["opt_idx"],
+                    metric=monitor_val,
+                )
                 self.scheduler_progress.increment_completed()
 
     def _get_monitor_value(self, key: str) -> Any:

pytorch_lightning/strategies/deepspeed.py

@@ -24,7 +24,6 @@
 import torch
 from torch.nn import Module
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
 
 import pytorch_lightning as pl
 from pytorch_lightning.core.optimizer import _get_default_scheduler_config, _init_optimizers_and_lr_schedulers
@@ -41,7 +40,7 @@
 from pytorch_lightning.utilities.imports import _DEEPSPEED_AVAILABLE
 from pytorch_lightning.utilities.model_helpers import is_overridden
 from pytorch_lightning.utilities.seed import reset_seed
-from pytorch_lightning.utilities.types import _PATH, LRSchedulerTypeTuple, STEP_OUTPUT
+from pytorch_lightning.utilities.types import _PATH, LRSchedulerConfig, LRSchedulerTypeUnion, STEP_OUTPUT
 from pytorch_lightning.utilities.warnings import rank_zero_warn, WarningCache
 
 warning_cache = WarningCache()
@@ -399,7 +398,7 @@ def _setup_model_and_optimizers(self, model: Module, optimizers: List[Optimizer]
         return self.model, [optimizer]
 
     def _setup_model_and_optimizer(
-        self, model: Module, optimizer: Optimizer, lr_scheduler: Optional[_LRScheduler] = None
+        self, model: Module, optimizer: Optimizer, lr_scheduler: Optional[LRSchedulerTypeUnion] = None
     ):
         """Initialize one model and one optimizer with an optional learning rate scheduler.
 
@@ -445,7 +444,7 @@ def init_deepspeed(self):
         else:
             self._initialize_deepspeed_inference(model)
 
-    def _init_optimizers(self) -> Tuple[Optimizer, Optional[Union[LRSchedulerTypeTuple]], Optional[int]]:
+    def _init_optimizers(self) -> Tuple[Optimizer, Optional[List[LRSchedulerConfig]], Optional[int]]:
         optimizers, schedulers, optimizer_frequencies = _init_optimizers_and_lr_schedulers(self.lightning_module)
         if len(optimizers) > 1 or len(schedulers) > 1:
             raise MisconfigurationException(

pytorch_lightning/strategies/horovod.py

@@ -17,7 +17,6 @@
 import torch
 import torch.nn as nn
 from torch.optim import Optimizer
-from torch.optim.lr_scheduler import _LRScheduler
 
 import pytorch_lightning as pl
 from pytorch_lightning.core.optimizer import LightningOptimizer
@@ -105,8 +104,7 @@ def _unpack_lightning_optimizer(opt):
         lr_schedulers = self.lightning_module.trainer.lr_schedulers
         for scheduler in lr_schedulers:
             scheduler = scheduler["scheduler"]
-            if isinstance(scheduler, _LRScheduler):
-                scheduler.base_lrs = [lr * self.world_size for lr in scheduler.base_lrs]
+            scheduler.base_lrs = [lr * self.world_size for lr in scheduler.base_lrs]
 
         # Horovod: broadcast parameters & optimizer state to ensure consistent initialization
         hvd.broadcast_parameters(self.lightning_module.state_dict(), root_rank=0)

pytorch_lightning/trainer/connectors/logger_connector/fx_validator.py

@@ -43,6 +43,7 @@ class _LogOptions(TypedDict):
         "optimizer_step": _LogOptions(
             allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
         ),
+        "lr_scheduler_step": None,
         "on_before_zero_grad": _LogOptions(
             allowed_on_step=(False, True), allowed_on_epoch=(False, True), default_on_step=True, default_on_epoch=False
         ),

pytorch_lightning/utilities/types.py

@@ -46,12 +46,11 @@
 EVAL_DATALOADERS = Union[DataLoader, Sequence[DataLoader]]
 
 
-# Copied from `torch.optim.lr_scheduler.pyi`
-# Missing attributes were added to improve typing
+# Inferred from `torch.optim.lr_scheduler.pyi`
 class _LRScheduler:
     optimizer: Optimizer
 
-    def __init__(self, optimizer: Optimizer, last_epoch: int = ...) -> None:
+    def __init__(self, optimizer: Optimizer, *args: Any, **kwargs: Any) -> None:
         ...
 
     def state_dict(self) -> dict:
@@ -60,13 +59,7 @@ def state_dict(self) -> dict:
     def load_state_dict(self, state_dict: dict) -> None:
         ...
 
-    def get_last_lr(self) -> List[float]:
-        ...
-
-    def get_lr(self) -> float:
-        ...
-
-    def step(self, epoch: Optional[int] = ...) -> None:
+    def step(self, *args: Any, **kwargs: Any) -> None:
         ...
 
 

tests/models/test_hooks.py

@@ -326,6 +326,17 @@ def _auto_train_batch(trainer, model, batches, device=torch.device("cpu"), curre
                         args=(current_epoch, i, ANY, 0, ANY),
                         kwargs=dict(on_tpu=False, using_lbfgs=False, using_native_amp=using_native_amp),
                     ),
+                    *(
+                        [
+                            dict(
+                                name="lr_scheduler_step",
+                                args=(ANY,),
+                                kwargs=dict(optimizer_idx=None, metric=None),
+                            )
+                        ]
+                        if i == (trainer.num_training_batches - 1)
+                        else []
+                    ),
                     dict(name="Callback.on_train_batch_end", args=(trainer, model, dict(loss=ANY), ANY, i)),
                     dict(name="on_train_batch_end", args=(dict(loss=ANY), ANY, i)),
                     dict(name="Callback.on_batch_end", args=(trainer, model)),

tests/trainer/logging_/test_logger_connector.py

@@ -233,6 +233,7 @@ def test_fx_validator_integration(tmpdir):
         "configure_callbacks": "You can't",
         "on_validation_model_eval": "You can't",
         "on_validation_model_train": "You can't",
+        "lr_scheduler_step": "You can't",
         "summarize": "not managed by the `Trainer",
     }
     model = HookedModel(not_supported)