bc

Author: rohitgr7

README.md

@@ -319,6 +319,7 @@ class LitAutoEncoder(pl.LightningModule):
         self.automatic_optimization = False
 
     def training_step(self, batch, batch_idx):
+        # access your optimizers with use_pl_optimizer=False. Default is True
         opt_a, opt_b = self.optimizers(use_pl_optimizer=True)
 
         loss_a = ...

pytorch_lightning/trainer/training_loop.py

@@ -838,14 +838,19 @@ def build_train_args(self, batch, batch_idx, opt_idx, hiddens):
         # enable not needing to add opt_idx to training_step
         args = [batch, batch_idx]
 
-        if len(self.trainer.optimizers) > 1 and self.automatic_optimization:
+        if len(self.trainer.optimizers) > 1:
             if self.trainer.has_arg("training_step", "optimizer_idx"):
+                if not self.automatic_optimization:
+                    self.warning_cache.warn(
+                        "`training_step` hook signature has changed in v1.3."
+                        " `optimizer_idx` argument has been removed in case of manual optimization. Support for"
+                        " the old signature will be removed in v1.5", DeprecationWarning
+                    )
                 args.append(opt_idx)
-            else:
-                num_opts = len(self.trainer.optimizers)
+            elif not self.trainer.has_arg("training_step", "optimizer_idx") and self.automatic_optimization:
                 raise ValueError(
-                    f"Your LightningModule defines {num_opts} optimizers but"
-                    f' training_step is missing the "optimizer_idx" argument.'
+                    f"Your LightningModule defines {len(self.trainer.optimizers)} optimizers but"
+                    ' `training_step` is missing the `optimizer_idx` argument.'
                 )
 
         # pass hiddens if using tbptt

tests/deprecated_api/test_remove_1-5.py

@@ -16,6 +16,7 @@
 from unittest import mock
 
 import pytest
+from torch import optim
 
 from pytorch_lightning import Callback, Trainer
 from pytorch_lightning.loggers import WandbLogger
@@ -74,3 +75,25 @@ def test_v1_5_0_running_sanity_check():
     trainer = Trainer()
     with pytest.deprecated_call(match='has been renamed to `Trainer.sanity_checking`'):
         assert not trainer.running_sanity_check
+
+
+def test_old_training_step_signature_with_opt_idx_manual_opt(tmpdir):
+
+    class OldSignatureModel(BoringModel):
+
+        def __init__(self):
+            super().__init__()
+            self.automatic_optimization = False
+
+        def training_step(self, batch, batch_idx, optimizer_idx):
+            assert optimizer_idx is not None
+            return super().training_step(batch, batch_idx)
+
+        def configure_optimizers(self):
+            return [optim.SGD(self.parameters(), lr=1e-2), optim.SGD(self.parameters(), lr=1e-2)]
+
+    model = OldSignatureModel()
+    trainer = Trainer(default_root_dir=tmpdir, fast_dev_run=2)
+
+    with pytest.deprecated_call(match="`training_step` .* `optimizer_idx` .* manual .* will be removed in v1.5"):
+        trainer.fit(model)

tests/trainer/optimization/test_multiple_optimizers.py

@@ -14,6 +14,7 @@
 """
 Tests to ensure that the behaviours related to multiple optimizers works
 """
+import pytest
 import torch
 
 import pytorch_lightning as pl
@@ -90,11 +91,6 @@ def training_epoch_end(self, outputs) -> None:
             # outputs should be an array with an entry per optimizer
             assert len(outputs) == 2
 
-        def configure_optimizers(self):
-            optimizer = torch.optim.SGD(self.layer.parameters(), lr=0.1)
-            optimizer_2 = torch.optim.SGD(self.layer.parameters(), lr=0.1)
-            return optimizer, optimizer_2
-
     model = TestModel()
     model.val_dataloader = None
 
@@ -154,3 +150,20 @@ def training_epoch_end(self, outputs) -> None:
     trainer.fit(model)
 
     assert model.training_step_called
+
+
+def test_multiple_optimizers_no_opt_idx_argument(tmpdir):
+    """
+    Test that an error is raised if no optimizer_idx is present when
+    multiple optimizeres are passed in case of automatic_optimization
+    """
+
+    class TestModel(MultiOptModel):
+
+        def training_step(self, batch, batch_idx):
+            return super().training_step(batch, batch_idx)
+
+    trainer = pl.Trainer(default_root_dir=tmpdir, fast_dev_run=2)
+
+    with pytest.raises(ValueError, match='`training_step` is missing the `optimizer_idx`'):
+        trainer.fit(TestModel())