Are the epoch_end hooks in the LightningModule needed?

[ananthsub]

Lightning has had some recent issues filed around these hooks:

• training_epoch_end
• validation_epoch_end
• test_epoch_end
• predict_epoch_end

Examples:

• outputs in training_epoch_end contain only outputs from last batch repeated #8603
• Add predict_epoch_end and predict_step_end methods to LightningModule #8657
• Ordering of hooks #8670
• Inconsistent API for on_predict_epoch_end #8479

Do we really need these on the LightningModule? Why don't users track these outputs directly inside the LightningModule?

Asking users to do this can offer some benefits:

• less API surface area for the trainer reduces risk of bugs like this. these bugs hurt user trust because the control flow isn't visible to end users, so removing this class of bugs can have a disproportionate benefit to user trust
• the trainer doesn't take responsibility of data it doesn't directly use. this simplifies the loop internals
• less "framework magic" means more readable user code because this tracking is explicit
• because the tracking is explicit, the responsibility of testing also falls to users
• we avoid API confusion: for instance, when should users implement something in training_epoch_end vs on_train_epoch_end ? this can improve the onboarding experience (one less class of hooks to learn about, only 1 way to do things).
• it can also improve performance: if users implement something in training_epoch_end and don't use outputs, the trainer accumulates results needlessly which wastes memory. This is not clearly visible to the user either.

Cons:

• (marginally) more boilerplate code in LightningModules. for instance, users would need to pay attention to resetting the results in on_{train/val}_epoch_start unless they want to accumulate results across epochs.

@PyTorchLightning/core-contributors

[s-rog]

I agree with removing the duplicated hooks, the con mentioned would not affect my workflow much.

[aazzolini]

I see as a fundamental property of programming that implementing an empty, non-recursive function should never by itself affect the memory complexity of a program. Here by the simple fact that we're overriding a method, our trainer memory complexity goes from O(1) to O(N) (N = batches / epoch).
Since

• there is no precedent that I know of of using a "is overriden" check for gating such large overhead.
• there is no limit on the size of an epoch -- and as a matter of fact we have various use-cases at Facebook that would have an output list greater than the memory.

I see this as a harmful thing to have, and I fully support deprecating these.

I believe that there could be more elegant and explicit ways of supporting this.
E.g. would it be possible to implement these as a callbacks?

[ananthsub]

@tchaton we can provide users code examples & tutorials demonstrating how they can store and post-process outputs, the same way we're doing now for callbacks. i find that easier than needing to explain when the users should implement training_epoch_end vs on_train_epoch_end, and why the two can have dramatically different effects on training speed

[maximsch2]

I agree with @ananthsub here, but think there is potentially even better solution for users looking for simpler things:

1. Let's focus on core LightningModel having only minimally required set of hooks that are orthogonal in functionality
2. For backwards-compat and/or examples on how to do things we can also provide SimpleLightningModel that will implement training_epoch_end hooks if we want to.

class SimpleLightningModule(LightningModule):
    def __init__(self):
           self.__train_outputs = []

    def on_train_epoch_end(self):
         self.training_epoch_end(self.__train_outputs)
         self.__train_outputs = []

    def training_step(self, *args):
         result = self.simple_training_step(*args)
         self.__train_outputs.append(result)
         return result
     
    def simple_training_step(self, *args):
          pass # To be implemented by a user

    def training_epoch_end(self, *args):
          pass # To be implemented by a user

(we might need a bit more magic if we want full backwards compat to see what people are overriding and how we need to wrap those functions to enable what we want instead of renaming them).

[ananthsub]

@maximsch2 - I filed this issue: #8731 - let's follow up there to keep the discussion on one place going forward

[carmocca]

let's follow up there to keep the discussion on one place going forward

Locking this one 🔒