[1/2] Deprecate outputs in on_train_epoch_end hooks

CHANGELOG.md

@@ -207,6 +207,8 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Deprecated
 
+- Deprecated `outputs` in both `LightningModule.on_train_epoch_end` and `Callback.on_train_epoch_end` hooks ([#7339](https://github.com/PyTorchLightning/pytorch-lightning/pull/7339))
+
 
 - Deprecated `Trainer.truncated_bptt_steps` in favor of `LightningModule.truncated_bptt_steps` ([#7323](https://github.com/PyTorchLightning/pytorch-lightning/pull/7323))
 
@@ -217,7 +219,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Deprecated the `save_function` property from the `ModelCheckpoint` callback ([#7201](https://github.com/PyTorchLightning/pytorch-lightning/pull/7201))
 
 
-- Deprecated `LightningModule.write_predictions` and `LigtningModule.write_predictions_dict` ([#7066](https://github.com/PyTorchLightning/pytorch-lightning/pull/7066))
+- Deprecated `LightningModule.write_predictions` and `LightningModule.write_predictions_dict` ([#7066](https://github.com/PyTorchLightning/pytorch-lightning/pull/7066))
 
 
 - Deprecated `TrainerLoggingMixin` in favor of a separate utilities module for metric handling ([#7180](https://github.com/PyTorchLightning/pytorch-lightning/pull/7180))

pytorch_lightning/accelerators/accelerator.py

@@ -28,7 +28,7 @@
 from pytorch_lightning.utilities import _NATIVE_AMP_AVAILABLE, rank_zero_warn
 from pytorch_lightning.utilities.apply_func import apply_to_collection, move_data_to_device
 from pytorch_lightning.utilities.enums import AMPType, GradClipAlgorithmType, LightningEnum
-from pytorch_lightning.utilities.types import EPOCH_OUTPUT, STEP_OUTPUT
+from pytorch_lightning.utilities.types import STEP_OUTPUT
 
 if _NATIVE_AMP_AVAILABLE:
     from torch.cuda.amp import GradScaler
@@ -354,12 +354,8 @@ def clip_gradients(
             model=self.model,
         )
 
-    def on_train_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
-        """Hook to do something on the end of an training epoch
-
-        Args:
-            outputs: the outputs of the training steps
-        """
+    def on_train_epoch_end(self) -> None:
+        """Hook to do something on the end of an training epoch."""
         pass
 
     def on_train_end(self) -> None:

pytorch_lightning/callbacks/base.py

@@ -98,7 +98,9 @@ def on_train_epoch_start(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningMo
         """Called when the train epoch begins."""
         pass
 
-    def on_train_epoch_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', outputs: EPOCH_OUTPUT) -> None:
+    def on_train_epoch_end(
+        self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule', unused: Optional = None
+    ) -> None:
         """Called when the train epoch ends."""
         pass
 

pytorch_lightning/callbacks/early_stopping.py

@@ -161,7 +161,7 @@ def _should_skip_check(self, trainer) -> bool:
         from pytorch_lightning.trainer.states import TrainerFn
         return trainer.state.fn != TrainerFn.FITTING or trainer.sanity_checking
 
-    def on_train_epoch_end(self, trainer, pl_module, outputs) -> None:
+    def on_train_epoch_end(self, trainer, pl_module) -> None:
         if not self._check_on_train_epoch_end or self._should_skip_check(trainer):
             return
         self._run_early_stopping_check(trainer)

pytorch_lightning/callbacks/pruning.py

@@ -373,7 +373,7 @@ def on_before_accelerator_backend_setup(self, trainer, pl_module: LightningModul
                 self._original_layers.setdefault(id_, {"data": deepcopy(module), "names": []})
                 self._original_layers[id_]["names"].append((i, name))
 
-    def on_train_epoch_end(self, trainer, pl_module: LightningModule, outputs):
+    def on_train_epoch_end(self, trainer, pl_module: LightningModule):
         current_epoch = trainer.current_epoch
         prune = self._apply_pruning(current_epoch) if isinstance(self._apply_pruning, Callable) else self._apply_pruning
         amount = self.amount(current_epoch) if isinstance(self.amount, Callable) else self.amount

pytorch_lightning/core/hooks.py

@@ -235,7 +235,7 @@ def on_train_epoch_start(self) -> None:
         Called in the training loop at the very beginning of the epoch.
         """
 
-    def on_train_epoch_end(self, outputs: EPOCH_OUTPUT) -> None:
+    def on_train_epoch_end(self, unused: Optional = None) -> None:
         """
         Called in the training loop at the very end of the epoch.
         """

pytorch_lightning/trainer/callback_hook.py

@@ -96,7 +96,15 @@ def on_train_epoch_end(self, outputs: EPOCH_OUTPUT):
             outputs: List of outputs on each ``train`` epoch
         """
         for callback in self.callbacks:
-            callback.on_train_epoch_end(self, self.lightning_module, outputs)
+            if is_param_in_hook_signature(callback.on_train_epoch_end, "outputs"):
+                warning_cache.warn(
+                    "The signature of `Callback.on_train_epoch_end` has changed in v1.3."
+                    " `outputs` parameter has been removed."
+                    " Support for the old signature will be removed in v1.5", DeprecationWarning
+                )
+                callback.on_train_epoch_end(self, self.lightning_module, outputs)
+            else:
+                callback.on_train_epoch_end(self, self.lightning_module)
 
     def on_validation_epoch_start(self):
         """Called when the epoch begins."""

pytorch_lightning/trainer/trainer.py

@@ -1211,6 +1211,11 @@ def _cache_logged_metrics(self):
             self.logger_connector.cache_logged_metrics()
 
     def call_hook(self, hook_name: str, *args, **kwargs) -> Any:
+        # Note this implementation is copy/pasted into the TrainLoop class in TrainLoop._on_train_epoch_end_hook
+        # This was done to manage the deprecation of an argument to on_train_epoch_end
+        # If making chnages to this function, ensure that those changes are also made to
+        # TrainLoop._on_train_epoch_end_hook
+
         # set hook_name to model + reset Result obj
         skip = self._reset_result_and_set_hook_fx_name(hook_name)
 

pytorch_lightning/trainer/training_loop.py

@@ -31,6 +31,7 @@
 from pytorch_lightning.utilities.grads import grad_norm
 from pytorch_lightning.utilities.model_helpers import is_overridden
 from pytorch_lightning.utilities.parsing import AttributeDict
+from pytorch_lightning.utilities.signature_utils import is_param_in_hook_signature
 from pytorch_lightning.utilities.warnings import WarningCache
 
 
@@ -197,16 +198,14 @@ def reset_train_val_dataloaders(self, model) -> None:
 
     def track_epoch_end_reduce_metrics(self, epoch_output, batch_end_outputs):
 
+        hook_overridden = self._should_add_batch_output_to_epoch_output()
+
         # track the outputs to reduce at the end of the epoch
         for opt_idx, opt_outputs in enumerate(batch_end_outputs):
             sample_output = opt_outputs[-1]
 
             # decide if we need to reduce at the end of the epoch automatically
             auto_reduce_tng_result = isinstance(sample_output, Result) and sample_output.should_reduce_on_epoch_end
-            hook_overridden = (
-                is_overridden("training_epoch_end", model=self.trainer.lightning_module)
-                or is_overridden("on_train_epoch_end", model=self.trainer.lightning_module)
-            )
 
             # only track when a) it needs to be autoreduced OR b) the user wants to manually reduce on epoch end
             if not (hook_overridden or auto_reduce_tng_result):
@@ -218,6 +217,22 @@ def track_epoch_end_reduce_metrics(self, epoch_output, batch_end_outputs):
 
             epoch_output[opt_idx].append(opt_outputs)
 
+    def _should_add_batch_output_to_epoch_output(self) -> bool:
+        # We add to the epoch outputs if
+        # 1. The model defines training_epoch_end OR
+        # 2. The model overrides on_train_epoch_end which has `outputs` in the signature
+        # TODO: in v1.5 this only needs to check if training_epoch_end is overridden
+        lightning_module = self.trainer.lightning_module
+        if is_overridden("training_epoch_end", model=lightning_module):
+            return True
+
+        if is_overridden("on_train_epoch_end", model=lightning_module):
+            model_hook_fx = getattr(lightning_module, "on_train_epoch_end")
+            if is_param_in_hook_signature(model_hook_fx, "outputs"):
+                return True
+
+        return False
+
     def get_optimizers_iterable(self, batch_idx=None):
         """
         Generates an iterable with (idx, optimizer) for each optimizer.
@@ -593,9 +608,51 @@ def on_train_epoch_end(self, epoch_output: List[List[List[Result]]]) -> None:
             self.trainer.logger_connector.cache_logged_metrics()
 
         # call train epoch end hooks
-        self.trainer.call_hook('on_train_epoch_end', processed_epoch_output)
+        self._on_train_epoch_end_hook(processed_epoch_output)
         self.trainer.call_hook('on_epoch_end')
 
+    def _on_train_epoch_end_hook(self, processed_epoch_output) -> None:
+        # We cannot rely on Trainer.call_hook because the signatures might be different across
+        # lightning module and callback
+        # As a result, we need to inspect if the module accepts `outputs` in `on_train_epoch_end`
+
+        # This implementation is copied from Trainer.call_hook
+        hook_name = "on_train_epoch_end"
+
+        # set hook_name to model + reset Result obj
+        skip = self.trainer._reset_result_and_set_hook_fx_name(hook_name)
+
+        # always profile hooks
+        with self.trainer.profiler.profile(hook_name):
+
+            # first call trainer hook
+            if hasattr(self.trainer, hook_name):
+                trainer_hook = getattr(self.trainer, hook_name)
+                trainer_hook(processed_epoch_output)
+
+            # next call hook in lightningModule
+            model_ref = self.trainer.lightning_module
+            if is_overridden(hook_name, model_ref):
+                hook_fx = getattr(model_ref, hook_name)
+                if is_param_in_hook_signature(hook_fx, "outputs"):
+                    self.warning_cache.warn(
+                        "The signature of `ModelHooks.on_train_epoch_end` has changed in v1.3."
+                        " `outputs` parameter has been deprecated."
+                        " Support for the old signature will be removed in v1.5", DeprecationWarning
+                    )
+                    model_ref.on_train_epoch_end(processed_epoch_output)
+                else:
+                    model_ref.on_train_epoch_end()
+
+            # if the PL module doesn't have the hook then call the accelerator
+            # used to auto-reduce things for the user with Results obj
+            elif hasattr(self.trainer.accelerator, hook_name):
+                accelerator_hook = getattr(self.trainer.accelerator, hook_name)
+                accelerator_hook()
+
+        if not skip:
+            self.trainer._cache_logged_metrics()
+
     def run_training_batch(self, batch, batch_idx, dataloader_idx):
         # track grad norms
         grad_norm_dic = {}

tests/callbacks/test_callback_hook_outputs.py

@@ -34,9 +34,6 @@ def on_validation_batch_end(self, trainer, pl_module, outputs, batch, batch_idx,
         def on_test_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx):
             assert 'x' in outputs
 
-        def on_train_epoch_end(self, trainer, pl_module, outputs):
-            assert len(outputs) == trainer.num_training_batches
-
     class TestModel(BoringModel):
 
         def on_train_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx: int) -> None:
@@ -48,7 +45,7 @@ def on_validation_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx
         def on_test_batch_end(self, outputs, batch, batch_idx: int, dataloader_idx: int) -> None:
             assert 'x' in outputs
 
-        def on_train_epoch_end(self, outputs) -> None:
+        def training_epoch_end(self, outputs) -> None:
             assert len(outputs) == self.trainer.num_training_batches
 
     model = TestModel()

tests/callbacks/test_finetuning_callback.py

@@ -27,7 +27,7 @@
 
 class TestBackboneFinetuningCallback(BackboneFinetuning):
 
-    def on_train_epoch_end(self, trainer, pl_module, outputs):
+    def on_train_epoch_end(self, trainer, pl_module):
         epoch = trainer.current_epoch
         if self.unfreeze_backbone_at_epoch <= epoch:
             optimizer = trainer.optimizers[0]

tests/deprecated_api/test_remove_1-5.py

@@ -216,6 +216,53 @@ def test_v1_5_0_model_checkpoint_period(tmpdir):
         ModelCheckpoint(dirpath=tmpdir, period=1)
 
 
+def test_v1_5_0_old_on_train_epoch_end(tmpdir):
+    callback_warning_cache.clear()
+
+    class OldSignature(Callback):
+
+        def on_train_epoch_end(self, trainer, pl_module, outputs):  # noqa
+            ...
+
+    class OldSignatureModel(BoringModel):
+
+        def on_train_epoch_end(self, outputs):  # noqa
+            ...
+
+    model = BoringModel()
+    trainer = Trainer(default_root_dir=tmpdir, max_epochs=1, callbacks=OldSignature())
+
+    with pytest.deprecated_call(match="old signature will be removed in v1.5"):
+        trainer.fit(model)
+
+    callback_warning_cache.clear()
+
+    model = OldSignatureModel()
+
+    with pytest.deprecated_call(match="old signature will be removed in v1.5"):
+        trainer.fit(model)
+
+    trainer.train_loop.warning_cache.clear()
+
+    class NewSignature(Callback):
+
+        def on_train_epoch_end(self, trainer, pl_module):
+            ...
+
+    trainer.callbacks = [NewSignature()]
+    with no_deprecated_call(match="`Callback.on_train_epoch_end` signature has changed in v1.3."):
+        trainer.fit(model)
+
+    class NewSignatureModel(BoringModel):
+
+        def on_train_epoch_end(self):
+            ...
+
+    model = NewSignatureModel()
+    with no_deprecated_call(match="`ModelHooks.on_train_epoch_end` signature has changed in v1.3."):
+        trainer.fit(model)
+
+
 def test_v1_5_0_old_on_validation_epoch_end(tmpdir):
     callback_warning_cache.clear()
 

tests/models/test_hooks.py

@@ -17,7 +17,7 @@
 import pytest
 import torch
 
-from pytorch_lightning import Callback, Trainer
+from pytorch_lightning import Trainer
 from tests.helpers import BoringDataModule, BoringModel, RandomDataset
 from tests.helpers.runif import RunIf
 
@@ -92,21 +92,17 @@ def training_epoch_end(self, outputs):
 def test_training_epoch_end_metrics_collection_on_override(tmpdir):
     """ Test that batch end metrics are collected when training_epoch_end is overridden at the end of an epoch. """
 
-    class LoggingCallback(Callback):
+    class OverriddenModel(BoringModel):
 
-        def on_train_epoch_start(self, trainer, pl_module):
+        def __init__(self):
+            super().__init__()
             self.len_outputs = 0
 
-        def on_train_epoch_end(self, trainer, pl_module, outputs):
-            self.len_outputs = len(outputs)
-
-    class OverriddenModel(BoringModel):
-
         def on_train_epoch_start(self):
             self.num_train_batches = 0
 
-        def training_epoch_end(self, outputs):  # Overridden
-            return
+        def training_epoch_end(self, outputs):
+            self.len_outputs = len(outputs)
 
         def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
             self.num_train_batches += 1
@@ -123,22 +119,14 @@ def on_train_batch_end(self, outputs, batch, batch_idx, dataloader_idx):
     not_overridden_model = NotOverriddenModel()
     not_overridden_model.training_epoch_end = None
 
-    callback = LoggingCallback()
     trainer = Trainer(
         max_epochs=1,
         default_root_dir=tmpdir,
         overfit_batches=2,
-        callbacks=[callback],
     )
 
     trainer.fit(overridden_model)
-    # outputs from on_train_batch_end should be accessible in on_train_epoch_end hook
-    # if training_epoch_end is overridden
-    assert callback.len_outputs == overridden_model.num_train_batches
-
-    trainer.fit(not_overridden_model)
-    # outputs from on_train_batch_end should be empty
-    assert callback.len_outputs == 0
+    assert overridden_model.len_outputs == overridden_model.num_train_batches
 
 
 @RunIf(min_gpus=1)
@@ -334,9 +322,9 @@ def on_train_epoch_start(self):
             self.called.append("on_train_epoch_start")
             super().on_train_epoch_start()
 
-        def on_train_epoch_end(self, outputs):
+        def on_train_epoch_end(self):
             self.called.append("on_train_epoch_end")
-            super().on_train_epoch_end(outputs)
+            super().on_train_epoch_end()
 
         def on_validation_start(self):
             self.called.append("on_validation_start")

tests/trainer/logging_/test_logger_connector.py

@@ -678,7 +678,7 @@ def _assert_epoch_end(self, stage):
             acc.reset.asset_not_called()
             ap.reset.assert_not_called()
 
-        def on_train_epoch_end(self, outputs):
+        def on_train_epoch_end(self):
             self._assert_epoch_end('train')
 
         def on_validation_epoch_end(self, outputs):

tests/trainer/logging_/test_train_loop_logging.py

@@ -599,7 +599,7 @@ def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx, data
             # with func = np.mean if on_epoch else func = np.max
             self.count += 1
 
-        def on_train_epoch_end(self, trainer, pl_module, outputs):
+        def on_train_epoch_end(self, trainer, pl_module):
             self.make_logging(
                 pl_module, 'on_train_epoch_end', 8, on_steps=[False], on_epochs=self.choices, prob_bars=self.choices
             )

tests/trainer/test_training_loop.py

@@ -155,11 +155,6 @@ def training_epoch_end(self, outputs):
             [HookedModel._check_output(output) for output in outputs]
             super().training_epoch_end(outputs)
 
-        def on_train_epoch_end(self, outputs):
-            assert len(outputs) == 2
-            [HookedModel._check_output(output) for output in outputs]
-            super().on_train_epoch_end(outputs)
-
     model = HookedModel()
 
     # fit model