[bugfix] Resolve bug with multiple optimizers and toggle. (#5574)

* fix toggle_optimizer

* update doc

* resolve bug

* update

* Update pytorch_lightning/core/lightning.py

Co-authored-by: Rohit Gupta <[email protected]>

* update on comments

* update on comments

* update

Co-authored-by: Rohit Gupta <[email protected]>

(cherry picked from commit c76cc23)

Author: tchaton

CHANGELOG.md

@@ -164,6 +164,9 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Fixed
 
+- Fixed `toggle_optimizer` to reset `requieres_grad` state  ([#5574](https://github.com/PyTorchLightning/pytorch-lightning/pull/5574))
+
+
 - Fixed distributed setting and `ddp_cpu` only with `num_processes>1` ([#5297](https://github.com/PyTorchLightning/pytorch-lightning/pull/5297))
 
 
@@ -213,7 +216,6 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Check environ before selecting a seed to prevent warning message ([#4743](https://github.com/PyTorchLightning/pytorch-lightning/pull/4743))
 - Fixed signature mismatch in `model_to_device` of `DDPCPUHPCAccelerator` ([#5505](https://github.com/PyTorchLightning/pytorch-lightning/pull/5505))
 
-
 ## [1.1.3] - 2021-01-05
 
 ### Added

pytorch_lightning/core/lightning.py

@@ -1190,17 +1190,47 @@ def toggle_optimizer(self, optimizer: Optimizer, optimizer_idx: int):
 
         Override for your own behavior
 
+        It works with ``untoggle_optimizer`` to make sure param_requires_grad_state is properly reset.
+
         Args:
-            optimizer:
-            optimizer_idx:
+            optimizer: Current optimizer used in training_loop
+            optimizer_idx: Current optimizer idx in training_loop
         """
-        # Todo: required argument `optimizer_idx` is not used
-        for param in self.parameters():
-            param.requires_grad = False
+        param_requires_grad_state = {}
+        # make sure current optimizer is latest to be iterated over.
+        optimizers = [opt for opt in self.optimizers(use_pl_optimizer=False) if opt != optimizer] + [optimizer]
+        num_optimizers = len(optimizers) - 1
+        for opt_idx, opt in enumerate(optimizers):
+            for group in opt.param_groups:
+                for param in group['params']:
+                    if num_optimizers == opt_idx:
+                        # If a param appears in 2 optimizers, revert `requires_grad` to before toggle.
+                        if param in param_requires_grad_state:
+                            param.requires_grad = param_requires_grad_state[param]
+                    else:
+                        # save requires_grad for later restoration
+                        param_requires_grad_state[param] = param.requires_grad
+                        param.requires_grad = False
+
+        self._param_requires_grad_state = param_requires_grad_state
+
+    def untoggle_optimizer(self, optimizer_idx: int):
+        """
+        .. note:: Only called when using multiple optimizers
 
-        for group in optimizer.param_groups:
-            for param in group['params']:
-                param.requires_grad = True
+        Override for your own behavior
+
+        Args:
+            optimizer_idx: Current optimizer idx in training_loop
+        """
+        for opt_idx, opt in enumerate(self.optimizers(use_pl_optimizer=False)):
+            if optimizer_idx != opt_idx:
+                for group in opt.param_groups:
+                    for param in group['params']:
+                        if param in self._param_requires_grad_state:
+                            param.requires_grad = self._param_requires_grad_state[param]
+        # save memory
+        del self._param_requires_grad_state
 
     def optimizer_step(
         self,

pytorch_lightning/trainer/training_loop.py

@@ -795,6 +795,10 @@ def training_step_and_backward(self, split_batch, batch_idx, opt_idx, optimizer,
                 if self.trainer.terminate_on_nan:
                     self.trainer.detect_nan_tensors(result.loss)
 
+                if len(self.trainer.optimizers) > 1:
+                    # revert back to previous state
+                    self.trainer.get_model().untoggle_optimizer(opt_idx)
+
         return result
 
     def backward(self, result, optimizer, opt_idx, *args, **kwargs):