5
$\begingroup$

This should be a simple question. But it is vague to me. What do we mean by optimizer.zero_grad(). Consider SGD as an example: $W^{t+1}= W^{t}- \lambda g_t$. Which one becomes zero for each batch. It should be $g_t$ and not $W^t$. Am I right? Overall, for any optimizer, does it mean all other parameters except for $W^t$ and $W^{t+1}$?

$\endgroup$

1 Answer 1

6
$\begingroup$

You are right, $g_t$ is what becomes zero when we invoke optimizer.zero_grad().

After invoking zero_grad(), we compute the forward pass and invoke loss.backward(), which populates again $g_t$. Finally, we invoke optimizer.step(), which updates the weights, i.e. $W^{t+1}=W^{t} - \lambda g_t$ in your example.

We need to invoke zero_grad() to prevent loss.backward() from accumulating the new gradient values with the ones from the previous step.

Note that, each optimizer has its own way to use the gradients to update the weights.

$\endgroup$
1
  • 2
    $\begingroup$ Also note that not calling zero_grad() on each iteration is a common strategy if the optimal batch size is larger than your device memory. It's called gradient accumulation. And forgetting to call it at all is a common mistake - that even ml luminaries like Andrej Karpathy make (I think in this video which he corrects - youtube.com/watch?v=VMj-3S1tku0) $\endgroup$ Nov 8, 2023 at 21:52

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge you have read our privacy policy.

Not the answer you're looking for? Browse other questions tagged or ask your own question.