Deprecate write_predictions on the LightningModule

[ananthsub]

What does this PR do?

Fixes #7065

Before submitting

• Was this discussed/approved via a GitHub issue? (not for typos and docs)
• Did you read the contributor guideline, Pull Request section?
• Did you make sure your PR does only one thing, instead of bundling different changes together?
• Did you make sure to update the documentation with your changes? (if necessary)
• Did you write any new necessary tests? (not for typos and docs)
• Did you verify new and existing tests pass locally with your changes?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is free to review the PR once the tests have passed.
Before you start reviewing make sure you have read Review guidelines. In short, see the following bullet-list:

• Is this pull request ready for review? (if not, please submit in draft mode)
• Check that all items from Before submitting are resolved
• Make sure the title is self-explanatory and the description concisely explains the PR
• Add labels and milestones (and optionally projects) to the PR so it can be classified

Did you have fun?

Make sure you had fun coding 🙃

[codecov]

Codecov Report

Merging #7066 (2151ba4) into master (591b9ce) will decrease coverage by 4%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #7066    +/-   ##
=======================================
- Coverage      91%     87%    -4%     
=======================================
  Files         198     198            
  Lines       12633   12635     +2     
=======================================
- Hits        11536   11011   -525     
- Misses       1097    1624   +527     

[tchaton]

Hey @ananthsub,

I was talking with @carmocca and we thought we might want to introduce a
Writer Callback to Lightning for users to save predictions on batch / epoch. It should work for all stages possibly and users could active them if needed.

We should add properly on_predict_step_end and on_predict_epoch_end hooks to the PredictLoop.

We also noticed some issues with trainer.predict(...) with spawn (ddp_spawn, tpu_spawn). Only rank 0 results are being returned to the main process. This would be an issue if those accelerators are used to serve models into endpoints.

Should we try to return all ranks results for predict ?

Best,
T.C

[ananthsub]

@tchaton I think it'd make sense to write independently from each rank instead of gathering all results on rank 0. In terms of returning results to users, couldn't the callback contain the return state, similar to how the model checkpoint callback tracks the top K model paths?

[awaelchli]

What do we recommend instead of write_predictions? The predict method of Trainer is only for when the model is already trained and needs a dataloader.

[carmocca]

@awaelchli The alternative is #7127

[awaelchli]

I have seen that but is it a replacement? write_predictions was a method that users could use anywhere in their LightningModule (whether we recommend that or not). I have not seen an option for this in #7127.

[ananthsub]

@awaelchli users can still use torch.save + file I/O from within the lightning module to save arbitrary values. I don't think there has to be a direct lightning API replacement in order to deprecate this - what do you think?

[tchaton]

LGTM !