02 Β· Optimizers βοΈ
The loss tells the model how wrong it is; the optimizer decides how to fix it. It's the algorithm that turns gradients into actual learning.
1. The 'Why'
After a backward pass, every parameter in your model has a gradient β a direction that would increase the loss. To learn, you want to move each parameter a little in the opposite direction, shrinking the loss. That sounds simple ("subtract a bit of the gradient from each weight"), and the most basic optimizer does exactly that. But doing it well is surprisingly subtle: step too far and training explodes; step too timidly and it crawls; treat every parameter identically and you ignore that some need bigger nudges than others. The optimizer is the algorithm that manages all of this, and choosing and configuring it is one of the highest-leverage decisions in training.
You met the raw mechanics in Module 02, where you manually wrote w -= 0.1 * w.grad inside a torch.no_grad() block. An optimizer packages that update β across all your model's parameters β into two clean method calls, and adds decades of research on top: momentum to smooth out noisy gradients, per-parameter adaptive step sizes, and more. The two you'll reach for constantly are SGD (simple, well-understood, the classic) and Adam (adaptive, forgiving, a great default). This sub-module explains what they do, what the all-important learning rate controls, and the non-negotiable zero_grad() β step() rhythm that makes the training loop correct.
2. Core Concepts
What an optimizer holds. You construct it with your model's parameters and a learning rate: optim.SGD(model.parameters(), lr=0.01). It keeps references to those parameters and updates them in place.
The learning rate (lr). The size of each update step β the single most important hyperparameter. Too high and training diverges; too low and it's painfully slow. Typical starting points: ~0.1β0.01 for SGD, ~1e-3 for Adam.
SGD vs. Adam. SGD (often with momentum) takes a straightforward step along the gradient; it's robust and generalizes well but can need careful lr tuning. Adam adapts the step size per parameter automatically, making it forgiving and fast to get working β a superb default.
The two-call rhythm. Every step does optimizer.zero_grad() (clear last step's accumulated gradients β remember Module 02!) then, after loss.backward(), optimizer.step() (apply the update). Order and placement matter.
3. Code in Action
Creating an optimizer
import torch
import torch.nn as nn
import torch.optim as optim
model = nn.Linear(4, 3)
# Classic stochastic gradient descent with momentum
sgd = optim.SGD(model.parameters(), lr=0.01, momentum=0.9)
# Adam β adaptive, a forgiving default (note the smaller lr)
adam = optim.Adam(model.parameters(), lr=1e-3)
The zero_grad β backward β step cycle
import torch
import torch.nn as nn
import torch.optim as optim
model = nn.Linear(4, 3)
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=1e-3)
x = torch.randn(8, 4)
y = torch.randint(0, 3, (8,))
optimizer.zero_grad() # 1. clear old gradients (they accumulate otherwise!)
logits = model(x) # 2. forward pass
loss = criterion(logits, y)# 3. compute the loss
loss.backward() # 4. backward pass -> fills every .grad
optimizer.step() # 5. update all parameters using those gradients
Watching a single parameter move
import torch
w = torch.tensor([0.0], requires_grad=True)
optimizer = torch.optim.SGD([w], lr=0.1) # optimize a lone parameter
for step in range(5):
optimizer.zero_grad()
loss = (w - 5) ** 2 # minimum is at w = 5
loss.backward()
optimizer.step() # each step nudges w toward 5
print(f"step {step}: w = {w.item():.3f}")
4. Common Pitfalls
Pitfall 1 β Forgetting optimizer.zero_grad(). Because gradients accumulate (Module 02), skipping this makes each step use the sum of all previous gradients β training destabilizes with no error message. Make zero_grad() the reflexive first line of every iteration.
Pitfall 2 β A learning rate that's wrong by an order of magnitude. If the loss explodes to NaN, your lr is almost certainly too high; if it barely moves over many epochs, it's too low. When a model "won't train," adjust the learning rate first, before anything fancier.
Pitfall 3 β Passing the wrong parameters (or a copy) to the optimizer. The optimizer must receive the actual model.parameters() you intend to train. Building the optimizer before moving the model to a device, or handing it parameters from a different model, means step() updates nothing useful. Construct the optimizer after the model is finalized and on its device.
5. Further Reading & Watch List
- π Deep Learning with PyTorch, Ch. 5 β "Optimizers a la carte"
- π Learn PyTorch β Setting up an optimizer
- π οΈ
torch.optimdocs Β·optim.AdamΒ·optim.SGD - βΆοΈ Karpathy β building an optimizer by hand
β¬ οΈ Prev: 01 Β· Loss Functions Β· β‘οΈ Next: 03 Β· The Training Loop