Support for PyTorch's optimize_for_inference mode

[ananthsub]

🚀 Feature

Leverage PyTorch's optimize_for_inference mode for performance benefits during model evaluation and inference

PyTorch has recently introduced an experimental API optimize_for_inference:

• https://pytorch.org/docs/master/generated/torch.inference_mode.html#torch.inference_mode
• https://pytorch.org/docs/master/notes/autograd.html#locally-disabling-gradient-computation
• Create a context manager to enable InferenceMode in python frontend pytorch/pytorch#56608 / Add inference mode python bindings and tests pytorch/pytorch#58045

Motivation

Reap performance improvements

Pitch

This can be used during Trainer.predict in place of the no_grad if optimize_for_inference is available: https://github.com/PyTorchLightning/pytorch-lightning/blob/4c79b3a5b343866217784c66d122819c59a92c1d/pytorch_lightning/trainer/trainer.py#L1078-L1083

Alternatives

Keep as is

Additional context

If you enjoy PL, check out our other projects:

• Metrics: Machine learning metrics for distributed, scalable PyTorch applications.
• Flash: The fastest way to get a Lightning baseline! A collection of tasks for fast prototyping, baselining, finetuning and solving problems with deep learning
• Bolts: Pretrained SOTA Deep Learning models, callbacks and more for research and production with PyTorch Lightning and PyTorch
• Lightning Transformers: Flexible interface for high performance research using SOTA Transformers leveraging Pytorch Lightning, Transformers, and Hydra.

[carmocca]

We could also enable it for test, right?

[ananthsub]

We could also enable it for test, right?

Yup! I wonder if we could with validation too, as long as we're careful to reset the mode before the next train loop call

[carmocca]

I wonder if we could with validation too, as long as we're careful to reset the mode before the next train loop call

Not sure about this one, I can imagine somebody saving a validation inference-only tensor and trying to use it on a computation that has a training tensor. This would fail, right?