Enable inference mode for testing and predicting

[tangbinh]

What does this PR do?

• As proposed in Support for PyTorch's optimize_for_inference mode #8499, we enable inference mode for evaluation and prediction. The feature was included recently in PyTorch 1.9 and comes with computational benefits.
• We also correct a link to the previous pull request Move logger and profiler finalization to trainer's teardown #8685 in the change log.

Does your PR introduce any breaking changes? If yes, please list them.

None

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 list all the breaking changes introduced by this pull request?
• Did you update the CHANGELOG? (not for typos, docs, test updates, or internal minor changes/refactorings)

PR review

Anyone in the community is welcome to review the PR.
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 #8813 (fe99e5f) into master (ca679cd) will decrease coverage by 4%.
The diff coverage is 100%.

@@           Coverage Diff           @@
##           master   #8813    +/-   ##
=======================================
- Coverage      92%     88%    -4%     
=======================================
  Files         179     179            
  Lines       14910   14908     -2     
=======================================
- Hits        13760   13145   -615     
- Misses       1150    1763   +613     

[ananthsub]

@tangbinh thanks for working on this! It'd be awesome to include information on performance speedups for some candidate models as well

[tangbinh]

Should we allow the users to opt out of the inference mode if they want to? I imagine some cases where the gradients are needed at test time (e.g. higher-order derivatives), and being able to turn it off by a flag might be important. I'm think of having _TORCH_GREATER_EQUAL_1_9 and os.getenv("NO_INFERENCE_MODE", "False").lower() == "false".

[tangbinh]

@tangbinh ... It'd be awesome to include information on performance speedups for some candidate models as well

Do you have suggestions for the candidate models and some datasets to test them on? I was running some simple tasks and didn't notice much of a difference.

[ananthsub]

Do you have suggestions for the candidate models and some datasets to test them on? I was running some simple tasks and didn't notice much of a difference.

@Borda do you recommend the models used for the parity benchmark?

[carmocca]

For the 2 cases where no_grad has been replaced already, we don't need to check the TrainerFn because by design, they will only be called during trainer.validate(), trainer.test(), or trainer.predict(). On these, no grad operations should take place so we are safe.

This is different from the following, which is the validation part of trainer.fit. This is where I'm not sure if no_grad should be kept:

https://github.com/PyTorchLightning/pytorch-lightning/blob/e0605472306d6b95bf2616ab88f8c29f4498402e/pytorch_lightning/loops/epoch/training_epoch_loop.py#L247

[tangbinh]

@ananthsub @tchaton @awaelchli @carmocca Please see this notebook for some performance tests. It looks like for some common models and datasets, enabling inference mode doesn't make much of a difference in terms of inference time. The observation is consistent on both GPU and CPU. That said, I think it doesn't hurt if we decide to move forward with inference mode.

[ananthsub]

thanks for working on this @tangbinh !

[tchaton]

Dear @tangbinh @ananthsub,

Could we get the benchmark for some traditional models to validate this doesn't bring a regression for Lightning performances + updated docs.

Best,
T.C

[tangbinh]

Could we get the benchmark for some traditional models to validate this doesn't bring a regression for Lightning performances + updated docs.

Would you mind proposing some traditional models that you have in mind? I've done some benchmark tests, although they don't follow previous examples as I couldn't find any good one.