Disable lr_scheduler.step() in manual optimization (#6825)

Co-authored-by: Adrian Wälchli <[email protected]>
Co-authored-by: thomas chaton <[email protected]>

Author: 3 people

CHANGELOG.md

@@ -96,6 +96,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added `model` parameter to precision plugins' `clip_gradients` signature ([#6764](https://github.com/PyTorchLightning/pytorch-lightning/pull/6764))
 
 
+- Added `is_last_batch` attribute to `Trainer` ([#6825](https://github.com/PyTorchLightning/pytorch-lightning/pull/6825))
+
+
 - Added `LightningModule.lr_schedulers()` for manual optimization  ([#6567](https://github.com/PyTorchLightning/pytorch-lightning/pull/6567))
 
 
@@ -129,6 +132,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Changed `PyTorchProfiler` to use `torch.autograd.profiler.record_function` to record functions ([#6349](https://github.com/PyTorchLightning/pytorch-lightning/pull/6349))
 
 
+- Disabled `lr_scheduler.step()` in manual optimization  ([#6825](https://github.com/PyTorchLightning/pytorch-lightning/pull/6825))
+
+
 - Changed warnings and recommendations for dataloaders in `ddp_spawn` ([#6762](https://github.com/PyTorchLightning/pytorch-lightning/pull/6762/))
 
 

docs/source/common/optimizers.rst

@@ -35,7 +35,7 @@ To manually optimize, do the following:
   * ``optimizer.step()`` to update your model parameters
 
 Here is a minimal example of manual optimization.
- 
+
 .. testcode:: python
 
     from pytorch_lightning import LightningModule

pytorch_lightning/trainer/connectors/optimizer_connector.py

@@ -32,7 +32,7 @@ def update_learning_rates(self, interval: str, monitor_metrics=None):
             interval: either 'epoch' or 'step'.
             monitor_metrics: dict of possible values to monitor
         """
-        if not self.trainer.lr_schedulers:
+        if not self.trainer.lr_schedulers or not self.trainer.train_loop.automatic_optimization:
             return
 
         for scheduler_idx, lr_scheduler in enumerate(self.trainer.lr_schedulers):

pytorch_lightning/trainer/optimizers.py

@@ -80,7 +80,8 @@ def init_optimizers(self, model: LightningModule) -> Tuple[List, List, List]:
                 ' * A list of the previously described dict format, with an optional "frequency" key (int)'
             )
 
-        lr_schedulers = self.configure_schedulers(lr_schedulers, monitor=monitor)
+        is_manual_optimization = not self.train_loop.automatic_optimization
+        lr_schedulers = self.configure_schedulers(lr_schedulers, monitor, is_manual_optimization)
         _validate_scheduler_optimizer(optimizers, lr_schedulers)
 
         return optimizers, lr_schedulers, optimizer_frequencies
@@ -98,8 +99,13 @@ def _convert_to_lightning_optimizer(trainer, optimizer):
             for opt_idx, opt in enumerate(self.optimizers)
         }
 
-    def configure_schedulers(self, schedulers: list, monitor: Optional[str] = None):
-        # Convert each scheduler into dict structure with relevant information
+    def configure_schedulers(
+        self,
+        schedulers: list,
+        monitor: Optional[str],
+        is_manual_optimization: bool,
+    ) -> List[Dict[str, Any]]:
+        """Convert each scheduler into dict structure with relevant information"""
         lr_schedulers = []
         default_config = _get_default_scheduler_config()
         for scheduler in schedulers:
@@ -117,6 +123,16 @@ def configure_schedulers(self, schedulers: list, monitor: Optional[str] = None):
                         f'The "interval" key in lr scheduler dict must be "step" or "epoch"'
                         f' but is "{scheduler["interval"]}"'
                     )
+                if is_manual_optimization:
+                    invalid_keys = {'interval', 'frequency', 'reduce_on_plateau', 'monitor', 'strict'}
+                    keys_to_warn = [k for k in scheduler.keys() if k in invalid_keys]
+
+                    if keys_to_warn:
+                        rank_zero_warn(
+                            f'The lr scheduler dict contains the key(s) {keys_to_warn}, but the keys will be ignored.'
+                            ' You need to call `lr_scheduler.step()` manually in manual optimization.',
+                            RuntimeWarning,
+                        )
 
                 scheduler['reduce_on_plateau'] = isinstance(
                     scheduler['scheduler'], optim.lr_scheduler.ReduceLROnPlateau

pytorch_lightning/trainer/training_loop.py

@@ -472,6 +472,7 @@ def run_training_epoch(self):
         for batch_idx, (batch, is_last_batch) in train_dataloader:
 
             self.trainer.batch_idx = batch_idx
+            self.trainer.is_last_batch = is_last_batch
 
             # ------------------------------------
             # TRAINING_STEP + TRAINING_STEP_END

tests/trainer/optimization/test_manual_optimization.py

@@ -1170,3 +1170,41 @@ def configure_optimizers(self):
     )
 
     trainer.fit(model)
+
+
+def test_lr_scheduler_step_not_called(tmpdir):
+    """
+    Test `lr_scheduler.step()` is not called in manual optimization.
+    """
+    class TestModel(BoringModel):
+        def __init__(self):
+            super().__init__()
+            self.automatic_optimization = False
+
+        def training_step(self, batch, batch_idx):
+            opt = self.optimizers()
+
+            output = self(batch)
+            loss = self.loss(batch, output)
+
+            opt.zero_grad()
+            self.manual_backward(loss)
+            opt.step()
+
+    model = TestModel()
+    model.training_step_end = None
+    model.training_epoch_end = None
+
+    trainer = Trainer(
+        max_epochs=1,
+        default_root_dir=tmpdir,
+        fast_dev_run=2,
+    )
+
+    with patch("torch.optim.lr_scheduler.StepLR.step") as lr_step:
+        trainer.fit(model)
+
+    # If a lr scheduler inherits `torch.optim.lr_scheduler._LRScheduler`,
+    # `.step()` is called once during its instantiation.
+    # Thus, the call count should be 1, not 0.
+    assert lr_step.call_count == 1

tests/trainer/optimization/test_optimizers.py

@@ -476,3 +476,25 @@ def configure_optimizers(self):
     trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
     with pytest.raises(MisconfigurationException, match="attatched with an optimizer that wasn't returned"):
         trainer.fit(model)
+
+
+def test_warn_invalid_scheduler_key_in_manual_optimization(tmpdir):
+    """
+    Test warning when invalid scheduler keys are provided in manual optimization.
+    """
+
+    class TestModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.automatic_optimization = False
+
+        def configure_optimizers(self):
+            opt = torch.optim.SGD(self.layer.parameters(), lr=0.1)
+            sch = torch.optim.lr_scheduler.StepLR(opt, step_size=1)
+            return [opt], [{"scheduler": sch, "interval": "epoch"}]
+
+    model = TestModel()
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=True)
+    with pytest.warns(RuntimeWarning, match='the keys will be ignored'):
+        trainer.fit(model)