Machine Learning
zero_grad()
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 0
The 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 separately
Transformers are usually trained with AdamW.
Exercise 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.
Show example 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))
Test yourself
-
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
Show answer
The gradients are wiped and the parameters aren't updated — step() skips parameters whose gradient is None.
-
Zero the gradients instead of setting them to None.
opt.zero_grad(____=False)
Show answer
opt.zero_grad(set_to_none=False)
-
Match the concepts.
- zero_grad()
- Sets gradients to None
- set_to_none=False
- A zero tensor
- weight_decay
- Pulls weights to zero
- AdamW
- Decoupled decay
Next lesson: Summary - Optimization
