Machine Learning
optimizer.step()
optimizer.step() updates weights with gradients.
optimizer.step()
step() updates each parameter according to its gradient. For SGD the rule is simple: p ← p - lr · p.grad. The optimiser takes the parameter list at construction:
opt = torch.optim.SGD(model.parameters(), lr=0.1)
w = torch.tensor([1.0], requires_grad=True)
opt2 = torch.optim.SGD([w], lr=0.1)
(3 * w).sum().backward() # grad = 3
opt2.step()
print(w.item()) # 1 - 0.1·3 = 0.7
Different layers can be given different learning rates with parameter groups: SGD([{"params": a, "lr": 0.01}, {"params": b}], lr=0.1).
Learning rate schedulers
Lowering the learning rate as training progresses usually gives a better result. torch.optim.lr_scheduler automates this; scheduler.step() is called at the end of every epoch:
opt = torch.optim.SGD(model.parameters(), lr=0.1)
zam = torch.optim.lr_scheduler.StepLR(opt, step_size=10, gamma=0.5)
for epoch in range(30):
egit_bir_epoch()
zam.step() # lr halves every 10 epochs
print(opt.param_groups[0]["lr"]) # 0.0125
The current rate is read in opt.param_groups[0]["lr"].
Exercise Step and scheduler
For w = torch.tensor([2.0], requires_grad=True), build SGD([w], lr=0.5), backpropagate the loss (w ** 2).sum() and call step(); print the new w with item(). Then attach StepLR(opt, step_size=2, gamma=0.1) to the same optimiser, advance the scheduler 4 times and print the learning rate with round(..., 4).
Show example solution
import torch
w = torch.tensor([2.0], requires_grad=True)
opt = torch.optim.SGD([w], lr=0.5)
(w ** 2).sum().backward()
opt.step()
print(w.item())
zam = torch.optim.lr_scheduler.StepLR(opt, step_size=2, gamma=0.1)
for _ in range(4):
opt.step()
zam.step()
print(round(opt.param_groups[0]["lr"], 4))
Test yourself
-
If w = 2, gradient = -4 and lr = 0.5, what is w after an SGD step?
Show answer
4 — 2 - 0.5·(-4) = 4.
-
Read the current learning rate.
lr = opt.____[0]["lr"]
Show answer
lr = opt.param_groups[0]["lr"]
-
Match the concepts.
- step()
- p ← p - lr·grad
- lr
- Step size
- StepLR
- Periodic decay
- param_groups
- Per-group setting
Next lesson: Summary - Training loop
