Disable lr_scheduler.step() in manual optimization

CHANGELOG.md

@@ -102,6 +102,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))
+
+
 ### Deprecated
 
 - `period` has been deprecated in favor of `every_n_val_epochs` in the `ModelCheckpoint` callback ([#6146](https://github.com/PyTorchLightning/pytorch-lightning/pull/6146))

docs/source/common/optimizers.rst

@@ -44,6 +44,8 @@ to manually manage the optimization process. To do so, do the following:
 
 .. warning:: Before 1.2, ``optimzer.step`` was calling ``optimizer.zero_grad()`` internally. From 1.2, it is left to the users expertise.
 
+.. warning:: Before 1.3, ``lr_scheduler.step`` was called automatically in both manual and automatic optimization. From 1.3, ``lr_scheduler.step`` is disabled in manual optimization so that you can call it at arbitrary intervals. Use ``self.lr_schedulers()`` in LightningModule to access your learning rate schedulers defined in ``LightningModule.configure_optimizers()``.
+
 .. tip:: To perform ``accumulate_grad_batches`` with one optimizer, you can do as such.
 
 .. tip:: ``self.optimizers()`` will return ``LightningOptimizer`` objects. You can access your own optimizer with ``optimizer.optimizer``. However, if you use your own optimizer to perform a step, Lightning won't be able to support accelerators and precision for you.
@@ -75,6 +77,7 @@ Here is the same example as above using a ``closure``.
 
  def training_step(self, batch, batch_idx):
  opt = self.optimizers()
+ sch = self.lr_schedulers()
 
  def closure():
  # Only zero_grad on the first batch to accumulate gradients
@@ -87,6 +90,7 @@ Here is the same example as above using a ``closure``.
  return loss
 
  opt.step(closure=closure)
+ sch.step()
 
 .. tip:: Be careful where you call ``zero_grad`` or your model won't converge. It is good pratice to call ``zero_grad`` before ``manual_backward``.
 

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 model.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,
+ ):
+ """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

tests/trainer/optimization/test_manual_optimization.py

@@ -1147,3 +1147,41 @@ def dis_closure():
 @RunIf(min_gpus=2, special=True)
 def test_step_with_optimizer_closure_with_different_frequencies_ddp_with_toggle_model(tmpdir):
  train_manual_optimization(tmpdir, "ddp", model_cls=TestManualOptimizationDDPModelToggleModel)
+
+
+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.
+ 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)