Use inference_mode instead of no_grad for pth v1.9.0

[vfdev-5]

🚀 Feature

The idea is to replace where it is appropriate torch.no_grad with inference_mode to speed-up computations: evalutation, metrics.

This works since pytorch v1.9.0.

• https://pytorch.org/docs/1.9.0/notes/autograd.html#inference-mode
• https://pytorch.org/docs/1.9.0/generated/torch.inference_mode.html?highlight=inference_mode#torch.inference_mode

Let's also benchmark the speed-up on a simple case.

[KickItLikeShika]

I will work on this, torch.no_grad exists in:

• Examples
• Docs
• Metrics
• Engine

I will start workin on them in this same order, what do you think?

[vfdev-5]

@KickItLikeShika it is more a discussion issue. Can you first benchmark how quicker inference_mode vs no_grad for one of our examples, e.g. cifar10 or just on metrics computation.

[KickItLikeShika]

Cool, will do that on CIFAR and share the results, maybe more complex use case like CIFAR, can get us more clear benchmarks.

[vfdev-5]

I agree. To see more clearly we can just benchmark metric updates

[KickItLikeShika]

Ran CIFAR10 example for 5 epochs each (no_grad, inference_mode), and averaged the time for evaluating test metrics. (on Colab GPU).

• torch.no_grad: 3.7 sec
• torch.inference_mode: 3.688 sec

NB: https://colab.research.google.com/drive/1zfWxg8H0XOutqXfh8BORXaz2sJ4hFE_-?usp=sharing

[vfdev-5]

More details on this issue and what to do to solve it:

• detailed benchmarking using metrics only. For example, using Accuracy metric together with a fake evaluation engine. is_completed method for metrics uses no_grad decorator.
  • once done, please report here the results and how to reproduce the results
• if we can gain ~15-20% propose a code change to make metrics computations faster for pytorch >= 1.9.0, without breaking the code for previous versions.