zero_grad()
- Machine Learning course
- Advanced
- Lesson 73/78
- 5 min
zero_grad() resets gradients at each step.
zero_grad()
opt.zero_grad() clears the gradients of all the parameters the optimiser manages. In PyTorch 2 the default is set_to_none=True: the gradients are set to None instead of a zero tensor; this saves a little memory and time:
opt.zero_grad()
print(w.grad) # None
opt.zero_grad(set_to_none=False)
# the gradient tensor is kept, its values become 0The position matters: zero_grad is called either at the start of the loop or right after step(); called between backward and step, it wipes the gradients and no update happens.
Weight decay
Optimisers can apply L2 regularisation directly: weight_decay pulls the weights slightly towards zero at every step. Because in Adam this penalty gets mixed with the adaptive scaling, AdamW applies the decay separately and usually generalises better:
opt = torch.optim.SGD(params, lr=0.1, weight_decay=1e-4) # g += wd·p
opt = torch.optim.AdamW(params, lr=1e-3, weight_decay=0.01) # p -= lr·wd·p separatelyTransformers are usually trained with AdamW.
Practice zero_grad in the wrong place
For w = 1.0, backpropagate the loss (w ** 2), then call zero_grad() and then step() (lr = 0.1). Print on one line that w didn't change (w.item()) and that the gradient is None. Then repeat in the right order (step first) and print the new w.
Starter code
import torch
w = torch.tensor(1.0, requires_grad=True)
opt = torch.optim.SGD([w], lr=0.1)
# buraya yazSample solution
import torch
w = torch.tensor(1.0, requires_grad=True)
opt = torch.optim.SGD([w], lr=0.1)
(w ** 2).backward()
opt.zero_grad()
opt.step()
print(w.item(), w.grad is None)
(w ** 2).backward()
opt.step()
opt.zero_grad()
print(round(w.item(), 3))Expected output
1.0 True 0.8
Questions and answers
Review the lesson questions together with their answers. To test yourself, solve the lesson interactively.
1 What happens if <code>zero_grad()</code> is called between <code>backward()</code> and <code>step()</code>?
- A normal update
- A double update
- An error
- The gradients are wiped and the parameters aren't updated Correct answer
Explanation: step() skips parameters whose gradient is None.
2 Zero the gradients instead of setting them to None.
opt.zero_grad(____=False)opt.zero_grad(set_to_none=False)Explanation: set_to_none=False keeps the gradient tensor.
3 Match the concepts.
| Term | Meaning |
|---|---|
zero_grad() |
Sets gradients to None |
set_to_none=False |
A zero tensor |
weight_decay |
Pulls weights to zero |
AdamW |
Decoupled decay |