During the implementation of a custom training loop in PyTorch, why is it necessary to call optimizer.zero_grad() at the start of each training iteration?
Select an answer to reveal the explanation.
Short Explanation
Here's the deal: PyTorch accumulates gradients by default, so you need to clear the previous step's gradients before the next backward pass. If you don't, the new gradients get added to the old ones and your updates go off track. Forgetting this is a classic mistake that catches people in production and on the exam.
Full Explanation
In PyTorch, the gradients of the model's parameters (weights and biases) are calculated during the backward pass (when loss.backward() is called) and stored in each parameter's .grad attribute. Crucially, PyTorch does not automatically overwrite these gradient values; instead, it accumulates (adds) them. This accumulation is highly useful for certain architectures, such as recurrent neural networks (RNNs) or when implementing virtual batch sizes that exceed hardware memory limits.
However, in standard mini-batch gradient descent, gradients must be computed independently for each individual mini-batch. If optimizer.zero_grad() is omitted, the gradients calculated for the current batch will be added to the gradients from all previous batches. As a result, the parameter updates performed by the optimizer (via optimizer.step()) will use an accumulation of gradients across multiple steps, causing the model to update parameters in incorrect directions and ultimately preventing convergence.
The correct choice is to clear the accumulated gradients from the previous step so they do not sum together during the new backward pass.
The learning-rate reset claim is incorrect because resetting the learning rate is not the purpose of zero_grad(). The learning rate is a hyperparameter and remains constant unless adjusted by a learning rate scheduler.
The weights-and-biases reset claim is incorrect because it describes resetting the model's parameters (weights and biases) to zero. Doing this would destroy all learned patterns and cause the model to fail to train, as neural networks must be initialized with non-zero (often randomized) weights to break symmetry.
The GPU-memory release claim is incorrect because it describes purging GPU memory. While memory management is important in deep learning, zero_grad() does not free GPU memory or delete tensors from the computational graph; it merely sets the values within existing gradient tensors to zero.