Background
Sections
IntroductionModule 01 β€” Tensors 🧊01 Β· Creating Tensors 🧊02 Β· Indexing & Reshaping πŸ”ͺ03 Β· Tensor Math βž—04 Β· Device Placement πŸ–₯️⚑Module 02 β€” Autograd βš™οΈ01 Β· The Computational Graph πŸ•ΈοΈ02 Β· Backward Pass & Gradients ⬅️03 Β· Turning Autograd Off πŸ›‘04 Β· Gradient Gotchas πŸͺ€Module 03 β€” Neural Networks01 Β· The `nn.Module` Basics 🧱02 Β· Common Layers 🧩03 Β· Building a Network πŸ—οΈ04 Β· Inspecting Models πŸ”Module 04 β€” Data Handling πŸ—‚οΈ01 Β· Dataset Basics πŸ“‡02 Β· The DataLoader 🚚03 Β· Transforms 🎨04 Β· Splits & Built-in Datasets βœ‚οΈModule 05 β€” The Training Loop πŸ”01 Β· Loss Functions 🎯02 Β· Optimizers βš™οΈ03 Β· The Training Loop πŸ”04 Β· Evaluation & Metrics πŸ“ŠModule 06 β€” Saving & Loading πŸ’Ύ01 Β· `state_dict` Basics πŸ’Ύ02 Β· Checkpoints & Resuming ⏸️03 Β· Loading for Inference πŸš€04 Β· Best Model & Early Stopping πŸ…Module 07 β€” Computer Vision πŸ‘οΈ01 Β· Convolutions πŸ”²02 Β· CNN Architecture πŸ›οΈ03 Β· Transfer Learning πŸ”04 Β· Image Classification Project πŸ§ͺModule 08 β€” NLP & Transformers πŸ’¬01 Β· Text Data & Tokenization πŸ”€02 Β· Embeddings 🧭03 Β· Recurrent Layers & LSTMs πŸ”„04 Β· Intro to Transformers ⚑Module 09 β€” Ecosystem: PyTorch Lightning ⚑01 Β· Why Lightning? πŸ€”02 Β· The LightningModule 🧩03 Β· The Trainer πŸŽ›οΈ04 Β· DataModules & Callbacks 🧰Module 10 β€” Deployment & Optimization πŸš€01 Β· Exporting Models πŸ“¦02 Β· `torch.compile` ⚑03 Β· Inference Optimization πŸͺΆ04 Β· Serving & Next Steps πŸŽ“

02 Β· Optimizers βš™οΈ

4 min read

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


⬅️ Prev: 01 Β· Loss Functions Β· ➑️ Next: 03 Β· The Training Loop