Remove duplicate epoch_end hooks in the Lightning Module

[ananthsub]

🚀 Feature

This isn't a feature as much as API refinement we should do.

These 2 sets of hooks are redundant:

• {training/validation/test}_epoch_end = defined inside the LightningModule: https://github.com/PyTorchLightning/pytorch-lightning/blob/b7a22ba046ba57072a71b12d16caff000e66f798/pytorch_lightning/core/lightning.py#L652-L691
• on_{train/val/test}_epoch_end = defined inside core/hooks: https://github.com/PyTorchLightning/pytorch-lightning/blob/b7a22ba046ba57072a71b12d16caff000e66f798/pytorch_lightning/core/hooks.py#L207-L211

Both take the step outputs from that epoch for post-processing, which is confusing for users. Which ones should users implement for post-processing?

@kaushikb11 @tchaton @edenafek

Motivation

Clarify the API

Pitch

Turn it back how it used to be before:

• {training/validation/test}_epoch_end gets the step outputs for post-processing. If the LightningModule implements this, then the Trainer retains the step outputs during execution. If not, we save memory and discard the step outputs.
• on_{train/val/test}_epoch_end takes no arguments and runs after {training/validation/test}_epoch_end is called. The user has a chance to persist outputs from the first hook before the second runs, but it doesn't need to be passed through by the trainer.

Alternatives

Additional context

[carmocca]

So, to recap, we currently have:

LightningModule:

• lightning.py: training_epoch_end(outputs)
• hooks.py: on_train_epoch_end(outputs)

Callback:

• base.py: on_train_epoch_end(outputs)

(Same thing for validation and test)

You are proposing removing the latter of the two options available for the LightningModule.

The disadvantage I see is the inconsistency with the Callback's on_train_end hook which will still take the outputs argument. A trade-off between redundancy and inconsistency.

cc: @PyTorchLightning/core-contributors for opinions

[ananthsub]

@carmocca

You are proposing removing the latter of the two options available for the LightningModule.

Not exactly. I'm proposing this:

LightningModule:

• lightning.py: training_epoch_end(outputs)
• hooks.py: on_train_epoch_end()

Callback:

• base.py: on_train_epoch_end(outputs)

Alternatively, we could add training_epoch_end(outputs) to the callback, and amend on_train_epoch_end to take no arguments to be consistent.

[carmocca]

Adding training_epoch_end(outputs) does not seem like an option to me, considering all callback hooks start with on_*.

On the other hand, I personally don't like having different arguments for the LightningModule on_train_epoch_end compared to the Callback's counterpart. You want the conversion between both to be as seamless as possible.

Another option would be to remove the LightningModule on_train_epoch_end entirely, but IMO it's nice to have to separate logic optional to the training procedure.

So yeah, no right or wrong answer, just opinions. Let's try to get some degree of consensus and then make a decision.

[ananthsub]

@PyTorchLightning/core-contributors what do you think?

👍 for this:
LightningModule:

• training_epoch_end(outputs)
• on_train_epoch_end()

Callback:

• on_train_epoch_end()

🚀 for this:
LightningModule:

• training_epoch_end(outputs) and remove on_train_epoch_end entirely

Callback:

• on_train_epoch_end(outputs)

😄 for this (keep as it is):

LightningModule:

• training_epoch_end(outputs)
• on_train_epoch_end(outputs)

Callback:

• on_train_epoch_end(outputs)

[ananthsub]

@carmocca beyond clarifying the API, another motivation is to save memory by avoiding storing the intermediary step outputs from the epoch. for large training runs, even just returning the loss can add up