PyTorch
no_grad
machine learning
deep learning
neural networks

What is the purpose of with torch.no_grad

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

`with torch.no_grad():` is a context manager in PyTorch that temporarily sets all the requires_grad flags to false, ensuring that the operations performed within this context do not compute gradients. The main purpose of using `torch.no_grad()` is to improve computational efficiency and resource utilization during specific operations that do not need gradient information, such as model inference.

Background on PyTorch Autograd

PyTorch is a popular open-source machine learning library that employs a dynamic computation graph, enabling runtime computation of gradients. When building machine learning models, especially neural networks, gradient computations are integral to the optimization process. This is achieved through PyTorch's autograd system, which automatically calculates the gradients of tensors. Typically, during training, autograd keeps track of operations on tensors that require gradients (i.e., those with `requires_grad=True`) to update model weights using backpropagation.

Technical Explanation

By default, PyTorch tracks all operations on tensors that have `requires_grad=True`. This tracking results in a complete computation graph that maps out dependencies between operations, which is later used to compute gradients during the backward pass. While this is essential for training, it is unnecessary for inference, where we simply forward propagate the data through the network to obtain predictions.

The `torch.no_grad()` context manager is used to suspend gradient tracking, thus ensuring that computational resources are not wasted on creating and storing unnecessary computation graphs. This contributes to faster execution times and reduced memory consumption during inference.

Example: Usage of `torch.no_grad()`

  • Backpropagation: It is crucial to switch off gradients with `torch.no_grad()` only when gradient information is unnecessary. If mistakenly used during training, it would disable the backpropagation process, preventing the model from learning.
  • Model Evaluation Mode: When evaluating a model, it is not only critical to disable gradients but also to call `.eval()` on the model to set the dropout and batch normalization layers to evaluation mode.
  • Automatic Mixed Precision (AMP): When using AMP for reducing memory and computational overhead via floating-point precision, `torch.no_grad()` should still be used as it complements AMP by further reducing unnecessary operations during inference.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.