Inconsistent API for on_predict_epoch_end

[ananthsub]

🚀 Feature

Background
We are auditing the Lightning components and APIs to assess opportunities for improvements:

• Review Lightning architecture & API #7740
• https://docs.google.com/document/d/1xHU7-iQSpp9KJTjI3As2EM0mfNHHr37WZYpDpwLkivA/edit#

One item that came up was on_predict_epoch_end() defined on the ModelHooks Mixin. This accepts an outputs: List[Any] argument. However, this is inconsistent with the other model hooks of the same type: on_train_epoch_end, on_validation_epoch_end, and on_test_epoch_end

Motivation

API consistency with other epoch end model hooks.

Pitch

• Add a prediction_epoch_end hook to the LightningModule
• Deprecate the outputs argument from on_predict_epoch_end in v1.5 and remove entirely in v1.7
• Update the prediction loop to only cache predictions if prediction_epoch_end is implemented

Users can optionally avoid this entirely and cache their prediction outputs as needed by implementing this logic in predict_step directly

Alternatives

Keep as is?

Additional context

cc @Borda @tchaton @justusschock @awaelchli @carmocca @ninginthecloud @daniellepintz @rohitgr7

[ishtos]

I'd like to work on this, could you assign this issue to me?

[mattcleigh]

So I just want to follow up on the standard way of now returning data using the predict step.
Ideally I would want this all defined in the model class so I dont have to tweak the prediction script for each model I use.

Would it be something like this, because it seems less intuitive to me compared to the older methods and adds alot of boilerplate code:

class Model(pl.LightningModule):
    def __init__(self, *args):
        super().__init__()
        self.save_hyperparameters(logger=False)
        self.predict_step_outputs = []
        self.predict_outputs = None
        ...

    def on_predict_epoch_start(self) -> None:
        self.predict_outputs = None

    def predict_step(self, batch: tuple, _batch_idx: int) -> None:
        outputs = ... # However to calculate the outputs
        self.predict_step_outputs.append(outputs)

    def on_predict_epoch_end(self, outputs) -> tuple:
        combined_predictions = ...  # However to combine, like T.vstack etc
        self.predict_outputs = combined_predictions
        self.predict_step_outputs.clear()

# Then in the export/predict script
trainer.predict(model=model, datamodule=datamodule)
outputs = model.predict_outputs

[carmocca]

That works. You could also do the combination after trainer.predict() so that you don't need the predict_outputs attribute. It's up to you now to do it however you prefer!

[mattcleigh]

So there is no way to overload how the trainer.predict() method combines the output of each batch? Will it always return as a list?

[carmocca]

Are you asking about the predictions returned? predictions = trainer.predict(..., return_predictions=True)?

[mattcleigh]

Exactly. By default it returns a list of whatever was returned by the model's predict_step for each batch pass. If I wanted to do some post processing on the entire collection of predictions, or some simple stacking so to get a single tensor out, I would have to do it afterwards. I was wondering if there was a built in function for this or would I have to define my own like:

predictions = trainer.predict(..., return_predictions=True)
predictions = model.process_and_combine(predictions)

[carmocca]

You would have to define your own