Make manual optimization mandatory for multiple optimizers

src/pytorch_lightning/callbacks/callback.py

@@ -241,7 +241,7 @@ def on_after_backward(self, trainer: "pl.Trainer", pl_module: "pl.LightningModul
         """Called after ``loss.backward()`` and before optimizers are stepped."""
 
     def on_before_optimizer_step(
-        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", optimizer: Optimizer, opt_idx: int
+        self, trainer: "pl.Trainer", pl_module: "pl.LightningModule", optimizer: Optimizer
     ) -> None:
         """Called before ``optimizer.step()``."""
 

src/pytorch_lightning/callbacks/finetuning.py

@@ -46,7 +46,7 @@ class BaseFinetuning(Callback):
         and should be used to freeze any modules parameters.
 
     ``finetune_function``: This method is called on every train epoch start and should be used to
-        ``unfreeze`` any parameters. Those parameters needs to be added in a new ``param_group``
+        ``unfreeze`` any parameters. Those parameters need to be added in a new ``param_group``
         within the optimizer.
 
     .. note:: Make sure to filter the parameters based on ``requires_grad``.
@@ -69,7 +69,7 @@ class BaseFinetuning(Callback):
         ...         # Here, we are freezing `feature_extractor`
         ...         self.freeze(pl_module.feature_extractor)
         ...
-        ...     def finetune_function(self, pl_module, current_epoch, optimizer, optimizer_idx):
+        ...     def finetune_function(self, pl_module, current_epoch, optimizer):
         ...         # When `current_epoch` is 10, feature_extractor will start training.
         ...         if current_epoch == self._unfreeze_at_epoch:
         ...             self.unfreeze_and_add_param_group(
@@ -290,18 +290,13 @@ def _store(
 
     def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -> None:
         """Called when the epoch begins."""
-        # import is here to avoid circular imports
-        from pytorch_lightning.loops.utilities import _get_active_optimizers
-
-        for opt_idx, optimizer in _get_active_optimizers(trainer.optimizers, trainer.optimizer_frequencies, 0):
+        for opt_idx, optimizer in enumerate(trainer.optimizers):
             num_param_groups = len(optimizer.param_groups)
-            self.finetune_function(pl_module, trainer.current_epoch, optimizer, opt_idx)
+            self.finetune_function(pl_module, trainer.current_epoch, optimizer)
             current_param_groups = optimizer.param_groups
             self._store(pl_module, opt_idx, num_param_groups, current_param_groups)
 
-    def finetune_function(
-        self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer, opt_idx: int
-    ) -> None:
+    def finetune_function(self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer) -> None:
         """Override to add your unfreeze logic."""
         raise NotImplementedError
 
@@ -389,9 +384,7 @@ def on_fit_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningModule") -
     def freeze_before_training(self, pl_module: "pl.LightningModule") -> None:
         self.freeze(pl_module.backbone)
 
-    def finetune_function(
-        self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer, opt_idx: int
-    ) -> None:
+    def finetune_function(self, pl_module: "pl.LightningModule", epoch: int, optimizer: Optimizer) -> None:
         """Called when the epoch begins."""
         if epoch == self.unfreeze_backbone_at_epoch:
             current_lr = optimizer.param_groups[0]["lr"]

src/pytorch_lightning/callbacks/stochastic_weight_avg.py

@@ -212,8 +212,8 @@ def on_train_epoch_start(self, trainer: "pl.Trainer", pl_module: "pl.LightningMo
                     "This may be caused by loading a checkpoint from an older version of PyTorch Lightning."
                 )
 
-            # We assert that there is only one optimizer on fit start, so know opt_idx is always 0
-            default_scheduler_cfg = LRSchedulerConfig(self._swa_scheduler, opt_idx=0)
+            # We assert that there is only one optimizer on fit start
+            default_scheduler_cfg = LRSchedulerConfig(self._swa_scheduler)
             assert default_scheduler_cfg.interval == "epoch" and default_scheduler_cfg.frequency == 1
 
             if trainer.lr_scheduler_configs:

src/pytorch_lightning/core/hooks.py

@@ -73,6 +73,7 @@ def on_train_batch_start(self, batch: Any, batch_idx: int) -> Optional[int]:
             batch_idx: the index of the batch
         """
 
+    # TODO: Should 'outputs' be renamed to 'output' (singular)?
     def on_train_batch_end(self, outputs: STEP_OUTPUT, batch: Any, batch_idx: int) -> None:
         """Called in the training loop after the batch.
 
@@ -229,7 +230,7 @@ def on_after_backward(self) -> None:
             Use the ``on_before_optimizer_step`` if you need the unscaled gradients.
         """
 
-    def on_before_optimizer_step(self, optimizer: Optimizer, optimizer_idx: int) -> None:
+    def on_before_optimizer_step(self, optimizer: Optimizer) -> None:
         """Called before ``optimizer.step()``.
 
         If using gradient accumulation, the hook is called once the gradients have been accumulated.
@@ -243,11 +244,10 @@ def on_before_optimizer_step(self, optimizer: Optimizer, optimizer_idx: int) ->
 
         Args:
             optimizer: Current optimizer being used.
-            optimizer_idx: Index of the current optimizer being used.
 
         Example::
 
-            def on_before_optimizer_step(self, optimizer, optimizer_idx):
+            def on_before_optimizer_step(self, optimizer):
                 # example to inspect gradient information in tensorboard
                 if self.trainer.global_step % 25 == 0:  # don't make the tf file huge
                     for k, v in self.named_parameters():