TopolearnThe optimizer and learning rate decide whether and how fast training converges.
An optimizer owns a list of parameters and updates them from their .grad. Create it once, before the loop:
opt = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=0.01)Each step is the same three calls: opt.zero_grad() clears old gradients, loss.backward() computes new ones, opt.step() applies the update.
Which optimizer?
SGD(lr, momentum=0.9): simple and strong for vision models with a tuned schedule.Adam: adaptive per-parameter step sizes, forgiving defaults (lr=1e-3).AdamW: Adam with decoupled weight decay; the default for transformers and most modern work.
The learning rate is the most important hyperparameter. Too high and the loss explodes or oscillates; too low and training crawls.
Optimizers can hold several parameter groups with different settings, for example no weight decay on biases and norms, or a smaller learning rate for a pretrained backbone:
opt = torch.optim.AdamW([
{"params": backbone.parameters(), "lr": 1e-5},
{"params": head.parameters()},
], lr=1e-3)Adam-style optimizers keep state (running averages) per parameter, roughly twice the model's size in memory, and that state must be saved with checkpoints to resume training exactly.
Example
import torch
w = torch.zeros(1, requires_grad=True)
b = torch.zeros(1, requires_grad=True)
opt = torch.optim.SGD([w, b], lr=0.1)
x = torch.linspace(-1, 1, 50)
y = 2 * x + 1 # the line we want to recover
for step in range(200):
loss = ((w * x + b - y) ** 2).mean()
opt.zero_grad() # 1. clear old gradients
loss.backward() # 2. compute new ones
opt.step() # 3. update w and b
print(round(w.item(), 3), round(b.item(), 3)) # 2.0 1.0Remember
- Loop:
zero_grad(),backward(),step() - AdamW is the common default; SGD with momentum for many vision tasks
- Parameter groups give different hyperparameters to different parts
Watch out for
- Creating the optimizer before replacing a layer (e.g. a new head), so the new parameters are never updated.