Improve manual optimization API

CHANGELOG.md

@@ -190,6 +190,10 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Refactored `EpochResultStore` ([#5522](https://github.com/PyTorchLightning/pytorch-lightning/pull/5522))
 
 
+- LightningOptimizer manual optimizer is more flexible and expose `toggle_model` ([#5771](https://github.com/PyTorchLightning/pytorch-lightning/pull/5771))
+
+
+
 ### Deprecated
 
 - Function `stat_scores_multiple_classes` is deprecated in favor of `stat_scores` ([#4839](https://github.com/PyTorchLightning/pytorch-lightning/pull/4839))

docs/source/common/optimizers.rst

@@ -21,46 +21,117 @@ Manual optimization
 For advanced research topics like reinforcement learning, sparse coding, or GAN research, it may be desirable
 to manually manage the optimization process. To do so, do the following:
 
-* Disable automatic optimization in Trainer:  Trainer(automatic_optimization=False)
+* Override your LightningModule ``automatic_optimization`` property to return ``False``
 * Drop or ignore the optimizer_idx argument
-* Use `self.manual_backward(loss)` instead of `loss.backward()` to automatically scale your loss
+* Use `self.manual_backward(loss)` instead of `loss.backward()`.
+
+.. note:: This is only recommended for experts who need ultimate flexibility. Lightning will handle only precision and accelerators logic. The users are left with zero_grad, accumulated_grad_batches, model toggling, etc..
+
+.. warning:: Before 1.2, ``optimzer.step`` was calling ``zero_grad`` internally. From 1.2, it is left to the users expertize.
+
+.. 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.
+
 
 .. code-block:: python
 
-    def training_step(self, batch, batch_idx, optimizer_idx):
+    def training_step(batch, batch_idx, optimizer_idx):
+        opt = self.optimizers()
+
+        loss = self.compute_loss(batch)
+        self.manual_backward(loss)
+        opt.step()
 
-        # 1. ignore optimizer_idx
-        # 2. `use_pl_optimizer=True` means `opt_g` and `opt_d` will be of type `LightningOptimizer`
-        # `LightningOptimizer` simply wrapped your optimizer and behave the same way !
-        # When calling `optimizer.step`, `LightningOptimizer` will just handle TPU, AMP, accumulate_grad_batches, etc ... for you.
+        # accumulate gradient batches
+        if batch_idx % 2 == 0:
+            opt.zero_grad()
 
-        # access your optimizers with `use_pl_optimizer=False` or `optimizer.optimizer` when using use_pl_optimizer=True
-        # use_pl_optimizer=True is the default
-        (opt_g, opt_d) = self.optimizers(use_pl_optimizer=True)
 
-        # do anything you want
-        loss_a = ...
+.. tip:: It is a good practice to provide the optimizer with a ``closure`` function that performs a ``forward`` and ``backward`` pass of your model. It is optional for most optimizers, but makes your code compatible if you switch to an optimizer which requires a closure.
 
-        # use self.backward which will also handle scaling the loss when using amp
-        self.manual_backward(loss_a, opt_g)
-        opt_g.step()
+Here is the same example as above using a ``closure``.
 
+.. code-block:: python
+
+    def training_step(batch, batch_idx, optimizer_idx):
+        opt = self.optimizers()
+
+        def forward_and_backward():
+            loss = self.compute_loss(batch)
+            self.manual_backward(loss)
 
-        # do anything you want
-        loss_b = ...
+        opt.step(closure=forward_and_backward)
 
-        # pass in any args that loss.backward() normally takes
-        self.manual_backward(loss_b, opt_d, retain_graph=True)
-        self.manual_backward(loss_b, opt_d)
-        opt_d.step()
+        # accumulate gradient batches
+        if batch_idx % 2 == 0:
+            opt.zero_grad()
+
+
+.. code-block:: python
 
-        # log losses
-        self.log('loss_a', loss_a)
-        self.log('loss_b', loss_b)
+    # Scenario for a GAN.
 
-.. note:: This is only recommended for experts who need ultimate flexibility
+    def training_step(...):
+        opt_gen, opt_dis = self.optimizers()
 
-Manual optimization does not yet support accumulated gradients but will be live in 1.1.0
+        # compute generator loss
+        loss_gen = self.compute_generator_loss(...)
+
+        # zero_grad needs to be called before backward
+        opt_gen.zero_grad()
+        self.manual_backward(loss_gen)
+        opt_gen.step()
+
+        # compute discriminator loss
+        loss_dis = self.compute_discriminator_loss(...)
+
+        # zero_grad needs to be called before backward
+        opt_dis.zero_grad()
+        self.manual_backward(loss_dis)
+        opt_dis.step()
+
+
+.. note:: ``LightningOptimizer`` provides a ``toggle_model`` function as a ``@context_manager`` for advanced users. It can be useful when performing gradient accumulation with several optimizers or training in a distributed setting.
+
+Here is an explanation of what it does:
+
+Considering the current optimizer as A and all other optimizers as B.
+Toggling means that all parameters from B exclusive to A will have their ``requires_grad`` attribute set to ``False``. Their original state will be restored when exiting the context manager.
+
+When performing gradient accumulation, there is no need to perform grad synchronization during the accumulation phase.
+Setting ``sync_grad`` to ``False`` will block this synchronization and improve your training speed.
+
+Here is an example on how to use it:
+
+.. code-block:: python
+
+
+    # Scenario for a GAN with gradient accumulation every 2 batches and optimized for multiple gpus.
+
+    def training_step(self, batch, batch_idx, ...):
+        opt_gen, opt_dis = self.optimizers()
+
+        accumulated_grad_batches = batch_idx % 2 == 0
+
+        # compute generator loss
+        def closure_gen():
+            loss_gen = self.compute_generator_loss(...)
+            self.manual_backward(loss_gen)
+            if accumulated_grad_batches:
+                opt_gen.zero_grad()
+
+        with opt_gen.toggle_model(sync_grad=accumulated_grad_batches):
+            opt_gen.step(closure=closure_gen)
+
+        def closure_dis():
+            loss_dis = self.compute_discriminator_loss(...)
+            self.manual_backward(loss_dis)
+            if accumulated_grad_batches:
+                opt_dis.zero_grad()
+
+        with opt_dis.toggle_model(sync_grad=accumulated_grad_batches):
+            opt_dis.step(closure=closure_dis)
 
 ------
 
@@ -166,7 +237,7 @@ returned as a dict which can contain the following keywords:
 * ``strict`` (optional): if set to ``True`` will enforce that value specified in ``monitor`` is available while trying
   to call ``scheduler.step()``, and stop training if not found. If ``False`` will only give a warning and continue training
   (without calling the scheduler).
-* ``name`` (optional): if using the :class:`~pytorch_lightning.callbacks.LearningRateMonitor` callback to monitor the 
+* ``name`` (optional): if using the :class:`~pytorch_lightning.callbacks.LearningRateMonitor` callback to monitor the
   learning rate progress, this keyword can be used to specify a specific name the learning rate should be logged as.
 
 .. testcode::
@@ -248,23 +319,6 @@ For example, here step optimizer A every 2 batches and optimizer B every 4 batch
             if batch_nb % 4 == 0 :
                optimizer.step(closure=closure)
 
-.. note:: When using ``Trainer(enable_pl_optimizer=True)``, ``.step`` accepts a boolean ``make_optimizer_step`` which can be used as follow.
-
-.. testcode::
-
-    def optimizer_zero_grad(self, current_epoch, batch_idx, optimizer, opt_idx):
-      optimizer.zero_grad()
-
-    # Alternating schedule for optimizer steps (ie: GANs)
-    def optimizer_step(self, current_epoch, batch_nb, optimizer, optimizer_idx, closure, on_tpu=False, using_native_amp=False, using_lbfgs=False):
-        # update generator opt every 2 steps
-        if optimizer_idx == 0:
-            optimizer.step(closure=closure, make_optimizer_step=(batch_nb % 2) == 0)
-
-        # update discriminator opt every 4 steps
-        if optimizer_idx == 1:
-            optimizer.step(closure=closure, make_optimizer_step=(batch_nb % 4) == 0)
-
 Here we add a learning-rate warm up
 
 .. testcode::

pytorch_lightning/accelerators/accelerator.py

@@ -238,7 +238,7 @@ def backward(
         self,
         closure_loss: torch.Tensor,
         optimizer: Optimizer,
-        opt_idx: int,
+        optimizer_idx: int,
         should_accumulate: bool,
         *args,
         **kwargs,
@@ -247,17 +247,15 @@ def backward(
 
         Args:
             closure_loss: a tensor holding the loss value to backpropagate
-            optimizer: the optimizer to do the step later on.
-            opt_idx: the index of the optimizer
             should_accumulate: whether to accumulate gradients
         """
-        self.training_type_plugin.pre_backward(closure_loss, should_accumulate, optimizer, opt_idx)
+        self.training_type_plugin.pre_backward(closure_loss, should_accumulate, optimizer, optimizer_idx)
 
         output = self.precision_plugin.backward(
-            self.lightning_module, closure_loss, optimizer, opt_idx, should_accumulate, *args, **kwargs
+            self.lightning_module, closure_loss, optimizer, optimizer_idx, should_accumulate, *args, **kwargs
         )
 
-        self.training_type_plugin.post_backward(closure_loss, should_accumulate, optimizer, opt_idx)
+        self.training_type_plugin.post_backward(closure_loss, should_accumulate, optimizer, optimizer_idx)
 
         return output
 

pytorch_lightning/core/lightning.py

@@ -1186,7 +1186,7 @@ def configure_optimizers(self):
         """
         rank_zero_warn("`configure_optimizers` must be implemented to be used with the Lightning Trainer")
 
-    def manual_backward(self, loss: Tensor, optimizer: Optimizer, *args, **kwargs) -> None:
+    def manual_backward(self, loss: Tensor, optimizer: Optional[Optimizer] = None, *args, **kwargs) -> None:
         """
         Call this directly from your training_step when doing optimizations manually.
         By using this we can ensure that all the proper scaling when using 16-bit etc has been done for you
@@ -1207,12 +1207,18 @@ def training_step(...):
                 self.manual_backward(loss, opt_a)
                 opt_a.step()
         """
+        if optimizer is not None:
+            rank_zero_warn(
+                "`optimizer` argument to `manual_backward` is deprecated in v1.2 and will be removed in v1.4",
+                DeprecationWarning
+            )
+
         # make sure we're using manual opt
         self._verify_is_manual_optimization('manual_backward')
 
         # backward
         self._running_manual_backward = True
-        self.trainer.train_loop.backward(loss, optimizer, -1, *args, **kwargs)
+        self.trainer.train_loop.backward(loss, optimizer=None, opt_idx=None, *args, **kwargs)
         self._running_manual_backward = False
 
     def backward(self, loss: Tensor, optimizer: Optimizer, optimizer_idx: int, *args, **kwargs) -> None: