Optimisation

Week 3 of 18 · Foundations · 7 days

Full curriculum
Week 03 · Foundations

Optimisation

Week 03 · Day 1 of 7

Optimisers

Why steepest descent is slow, and what momentum and Adam do about it

By 1281 words

Plain gradient descent takes a step straight downhill. That is the obvious thing to do and it is slow, for a reason that is easier to see than to describe.

The problem with steepest descent

import torch

# A bowl 200 times steeper in one direction than the other. The steep
# direction caps the learning rate at about 2/200, and the shallow one
# then needs an enormous number of those small steps.
def loss(p):
return 0.5 * (200 * p[0] ** 2 + p[1] ** 2)

p = torch.tensor([1.0, 1.0], requires_grad=True)
lr = 0.009
print('%6s %12s %12s %12s' % ('step', 'steep x', 'shallow y', 'loss'))
for step in range(151):
if p.grad is not None:
p.grad.zero_()
value = loss(p)
value.backward()
with torch.no_grad():
p -= lr * p.grad
if step in (0, 5, 20, 60, 150):
print('%6d %12.4f %12.4f %12.5f'
% (step, p[0].item(), p[1].item(), value.item()))
step steep x shallow y loss
0 -0.8000 0.9910 100.50000
5 0.2621 0.9472 11.19419
20 -0.0092 0.8271 0.36156
60 -0.0000 0.5761 0.16897
150 -0.0000 0.2553 0.03319

The steep direction is finished within a handful of steps. The shallow one is still most of the way from where it started after a hundred and fifty, because the rate that keeps the steep direction stable is far too small for it. Raising the rate makes the steep direction oscillate and then diverge, so you cannot. That is the tension every optimiser after this one is trying to resolve, and it gets worse as models grow, because the gap between the steepest and shallowest directions grows with them.

Momentum

Momentum: Keep a running average of recent gradients and step along that instead of along the current gradient. Directions that keep pointing the same way build up speed; directions that reverse every step cancel themselves out.
import torch

def loss(p):
return 0.5 * (200 * p[0] ** 2 + p[1] ** 2)

def run(beta, lr, steps):
p = torch.tensor([1.0, 1.0], requires_grad=True)
velocity = torch.zeros(2)
for _ in range(steps):
if p.grad is not None:
p.grad.zero_()
loss(p).backward()
with torch.no_grad():
velocity = beta * velocity + p.grad
p -= lr * velocity
return loss(p).item()

print('%-16s %12s %12s %12s' % ('', '20 steps', '60 steps', '150 steps'))
for label, beta in [('plain SGD', 0.0), ('momentum 0.9', 0.9)]:
print('%-16s %12.5f %12.5f %12.5f'
% (label, run(beta, 0.009, 20), run(beta, 0.009, 60),
run(beta, 0.009, 150)))
20 steps 60 steps 150 steps
plain SGD 0.36156 0.16897 0.03319
momentum 0.9 19.74603 0.18695 0.00001

Momentum looks worse before it looks better

At twenty steps momentum is around fifty times worse. A running average with beta = 0.9 settles at roughly ten times the size of a single gradient, so the effective step is ten times larger and it overshoots badly at the start. By a hundred and fifty steps it is thousands of times better than plain descent.

Two things follow from that. Momentum earns its place on long runs and can genuinely hurt on short ones, and if you add it without lowering the learning rate you should expect the early part of training to look wrong. Judging an optimiser on its first few hundred steps is how people talk themselves out of momentum.

Adam

Adam: Momentum on the gradient, and a second running average of the squared gradient used to divide the step. A parameter whose gradients are consistently large gets a smaller effective step and one whose gradients are tiny gets a larger one, so the two directions in that stretched bowl are treated on equal terms.
import torch

def loss(p):
return 0.5 * (20 * p[0] ** 2 + p[1] ** 2)

p = torch.tensor([1.0, 1.0], requires_grad=True)
m = torch.zeros(2)
v = torch.zeros(2)
b1, b2, eps, lr = 0.9, 0.999, 1e-8, 0.1

print('%6s %10s %10s %12s' % ('step', 'x', 'y', 'loss'))
for t in range(1, 10):
if p.grad is not None:
p.grad.zero_()
value = loss(p)
value.backward()
with torch.no_grad():
m = b1 * m + (1 - b1) * p.grad
v = b2 * v + (1 - b2) * p.grad ** 2
m_hat = m / (1 - b1 ** t) # bias correction
v_hat = v / (1 - b2 ** t)
p -= lr * m_hat / (v_hat.sqrt() + eps)
if t % 2 == 1:
print('%6d %10.4f %10.4f %12.5f'
% (t, p[0].item(), p[1].item(), value.item()))
step x y loss
1 0.9000 0.9000 10.50000
3 0.7016 0.7016 6.72693
5 0.5080 0.5080 3.82980
7 0.3234 0.3234 1.80171
9 0.1536 0.1536 0.58612

Look at the two coordinates: they are identical at every step, on a bowl where one direction is two hundred times steeper than the other. Dividing by the square root of the running average of squared gradients has removed the scale difference entirely, so the optimiser now walks diagonally to the minimum instead of zig-zagging down one axis and crawling along the other. That is the whole of what adaptive means.

Why the bias correction is there

Both running averages start at zero, so for the first few steps they are biased towards zero and the step size would be far too small. Dividing by 1 - beta**t undoes exactly that, and the correction fades away as t grows. Leave it out and the first hundred steps of training barely move, which is a bug people reinvent regularly when writing their own optimiser.

On a real model

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
def build(seed=0):
torch.manual_seed(seed)
return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))

def run(opt_fn, epochs=3, batch_size=64, seed=0, sched_fn=None):
torch.manual_seed(seed)
model = build(seed)
opt = opt_fn(model.parameters())
sched = sched_fn(opt) if sched_fn else None
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val = DataLoader(val_set, batch_size=512)
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
if sched is not None:
sched.step()
model.eval()
right = seen = 0
total = 0.0
with torch.no_grad():
for xb, yb in val:
out = model(xb)
total += loss_fn(out, yb).item() * yb.numel()
right += (out.argmax(1) == yb).sum().item()
seen += yb.numel()
return total / seen, right / seen
print('%-14s %12s %12s' % ('optimiser', 'val loss', 'val acc'))
for name, fn in [
('SGD 0.1', lambda p: torch.optim.SGD(p, lr=0.1)),
('SGD+mom 0.1', lambda p: torch.optim.SGD(p, lr=0.1, momentum=0.9)),
('Adam 1e-3', lambda p: torch.optim.Adam(p, lr=1e-3)),
('AdamW 1e-3', lambda p: torch.optim.AdamW(p, lr=1e-3)),
('RMSprop 1e-3', lambda p: torch.optim.RMSprop(p, lr=1e-3))]:
loss, acc = run(fn)
print('%-14s %12.4f %12.4f' % (name, loss, acc))
optimiser val loss val acc
SGD 0.1 0.2639 0.9235
SGD+mom 0.1 0.3766 0.9125
Adam 1e-3 0.2715 0.9230
AdamW 1e-3 0.2682 0.9220
RMSprop 1e-3 0.2355 0.9315

Notice that nothing wins convincingly

Five optimisers, and the spread between the best and the worst is about two points of accuracy, with plain SGD in the middle and momentum at the same rate behind it for the reason given above. On a small model and an easy dataset that is the normal result: the optimiser is not usually where your accuracy comes from.

Day 6 measures how much this model varies between random seeds, and the answer makes most of this table look smaller still. AdamW is the sensible default because it needs the least tuning to get somewhere reasonable, not because it wins races.

Day 1 takeaway

Steepest descent oscillates in steep directions and crawls in shallow ones. Momentum averages away the oscillation but multiplies the effective step, so lower the rate when you add it. Adam adds a per-parameter scale, which mostly buys insensitivity to your choices rather than a better final answer.
Week 03 · Day 2 of 7

Learning Rates and Schedules

Finding the rate with a range test, then decaying it properly

By 997 words

The learning rate matters more than the choice of optimiser, and it is the one hyperparameter you can find quickly rather than guess.

The range test

Learning rate range test: Train for a few hundred steps while increasing the learning rate exponentially, and record the loss at each rate. The loss falls, flattens, then explodes. The useful rate is roughly an order of magnitude below where it explodes, near the steepest part of the fall.
import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
import math
torch.manual_seed(0)

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))
opt = torch.optim.SGD(model.parameters(), lr=1e-5)
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=64, shuffle=True)

lo, hi, steps = 1e-5, 10.0, 120
gamma = (hi / lo) ** (1 / steps)
rates, losses = [], []
it = iter(loader)
for step in range(steps):
try:
xb, yb = next(it)
except StopIteration:
it = iter(loader)
xb, yb = next(it)
lr = lo * gamma ** step
for group in opt.param_groups:
group['lr'] = lr
opt.zero_grad()
loss = loss_fn(model(xb), yb)
loss.backward()
opt.step()
rates.append(lr)
losses.append(loss.item())

print('%12s %12s' % ('rate', 'loss'))
for i in range(0, steps, 12):
print('%12.5f %12.4f' % (rates[i], losses[i]))

best = min(range(len(losses)), key=lambda i: losses[i])
print('\nlowest loss at rate %.4f' % rates[best])
print('a sensible starting point is a little below that')
rate loss
0.00001 2.2856
0.00004 2.2757
0.00016 2.2586
0.00063 2.3156
0.00251 2.2612
0.01000 2.1424
0.03981 1.8005
0.15849 1.3027
0.63096 1.4018
2.51189 56.6162

lowest loss at rate 0.2818
a sensible starting point is a little below that

Schedules

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
torch.manual_seed(0)
model = nn.Sequential(nn.Linear(4, 4))
opt = torch.optim.SGD(model.parameters(), lr=0.1)

schedules = {
'step (x0.1 every 8)':
torch.optim.lr_scheduler.StepLR(opt, step_size=8, gamma=0.1),
}
print('%6s %14s %14s %14s' % ('epoch', 'step', 'cosine', 'one cycle'))

def trace(make, epochs=20):
o = torch.optim.SGD([torch.zeros(1, requires_grad=True)], lr=0.1)
s = make(o)
out = []
for _ in range(epochs):
out.append(o.param_groups[0]['lr'])
o.step()
s.step()
return out

step_lr = trace(lambda o: torch.optim.lr_scheduler.StepLR(o, 8, 0.1))
cosine = trace(lambda o: torch.optim.lr_scheduler.CosineAnnealingLR(o, 20))
onecycle = trace(lambda o: torch.optim.lr_scheduler.OneCycleLR(
o, max_lr=0.4, total_steps=20))
for e in range(0, 20, 2):
print('%6d %14.5f %14.5f %14.5f'
% (e, step_lr[e], cosine[e], onecycle[e]))
epoch step cosine one cycle
0 0.10000 0.10000 0.01600
2 0.10000 0.09755 0.14867
4 0.10000 0.09045 0.36333
6 0.10000 0.07939 0.39499
8 0.01000 0.06545 0.35637
10 0.01000 0.05000 0.28678
12 0.01000 0.03455 0.20000
14 0.01000 0.02061 0.11322
16 0.00100 0.00955 0.04364
18 0.00100 0.00245 0.00502
ScheduleShapeUse when
ConstantFlatA quick experiment, or Adam on a small problem
StepDrops by a factor at fixed epochsYou know roughly how long training takes
CosineSmooth decay to nearly zeroThe default for most modern training runs
One cycleWarms up, then decays below the startShort budgets; often the fastest to a good result
Reduce on plateauDrops when validation stallsYou cannot predict the length of the run
import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
def build(seed=0):
torch.manual_seed(seed)
return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))

def run(opt_fn, epochs=3, batch_size=64, seed=0, sched_fn=None):
torch.manual_seed(seed)
model = build(seed)
opt = opt_fn(model.parameters())
sched = sched_fn(opt) if sched_fn else None
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val = DataLoader(val_set, batch_size=512)
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
if sched is not None:
sched.step()
model.eval()
right = seen = 0
total = 0.0
with torch.no_grad():
for xb, yb in val:
out = model(xb)
total += loss_fn(out, yb).item() * yb.numel()
right += (out.argmax(1) == yb).sum().item()
seen += yb.numel()
return total / seen, right / seen
print('%-24s %12s %12s' % ('', 'val loss', 'val acc'))
for name, sched in [
('constant 0.1', None),
('cosine to zero',
lambda o: torch.optim.lr_scheduler.CosineAnnealingLR(o, 375)),
('one cycle, max 0.1',
lambda o: torch.optim.lr_scheduler.OneCycleLR(
o, max_lr=0.1, total_steps=375))]:
loss, acc = run(lambda p: torch.optim.SGD(p, lr=0.1, momentum=0.9),
epochs=3, sched_fn=sched)
print('%-24s %12.4f %12.4f' % (name, loss, acc))
val loss val acc
constant 0.1 0.3766 0.9125
cosine to zero 0.2330 0.9355
one cycle, max 0.1 0.2313 0.9400

Warmup

Large models, and anything using Adam with a large batch, often start with a few hundred steps at a very small rate that grows to the target. The reason is that Adam's running averages are unreliable in the first few steps, so a full-size step taken on that estimate can move the weights somewhere the model never recovers from.

import torch

def warmup_cosine(step, warmup=100, total=1000, peak=1.0):
if step < warmup:
return peak * step / warmup
progress = (step - warmup) / (total - warmup)
return peak * 0.5 * (1 + torch.cos(torch.tensor(3.14159 * progress)))

print('%8s %12s' % ('step', 'multiplier'))
for step in [0, 25, 50, 100, 300, 600, 900, 1000]:
print('%8d %12.4f' % (step, float(warmup_cosine(step))))
step multiplier
0 0.0000
25 0.2500
50 0.5000
100 1.0000
300 0.8830
600 0.4132
900 0.0302
1000 0.0000

Day 2 takeaway

Find the learning rate with a range test rather than guessing. Decay it over the run; cosine is a good default and one cycle is often faster on a short budget. Add a few hundred steps of warmup whenever you use Adam on a large model.
Week 03 · Day 3 of 7

Batch Size, Accumulation and Clipping

Gradient noise, fitting a large batch in a small memory, and the norm

By 991 words

Batch size looks like a memory setting and behaves like a hyperparameter, because it changes how noisy each gradient is and therefore how large a step you can safely take.

What the batch size actually changes

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))
loss_fn = nn.CrossEntropyLoss()

print('%8s %14s %16s' % ('batch', 'gradient norm', 'std across draws'))
for bs in [8, 32, 128, 512]:
norms = []
loader = DataLoader(train_set, batch_size=bs, shuffle=True)
it = iter(loader)
for _ in range(12):
xb, yb = next(it)
model.zero_grad()
loss_fn(model(xb), yb).backward()
flat = torch.cat([p.grad.flatten() for p in model.parameters()])
norms.append(flat.norm().item())
t = torch.tensor(norms)
print('%8d %14.4f %16.4f' % (bs, t.mean().item(), t.std().item()))
batch gradient norm std across draws
8 4.3953 0.5277
32 2.6351 0.2021
128 1.7999 0.1236
512 1.5807 0.0809

A larger batch gives a less noisy estimate of the true gradient. That is why it tolerates a larger learning rate, and it is the origin of the rule of thumb that doubling the batch size lets you roughly double the rate. The table below tests that rule, and finds where it stops being true.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
def build(seed=0):
torch.manual_seed(seed)
return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))

def run(opt_fn, epochs=3, batch_size=64, seed=0, sched_fn=None):
torch.manual_seed(seed)
model = build(seed)
opt = opt_fn(model.parameters())
sched = sched_fn(opt) if sched_fn else None
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val = DataLoader(val_set, batch_size=512)
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
if sched is not None:
sched.step()
model.eval()
right = seen = 0
total = 0.0
with torch.no_grad():
for xb, yb in val:
out = model(xb)
total += loss_fn(out, yb).item() * yb.numel()
right += (out.argmax(1) == yb).sum().item()
seen += yb.numel()
return total / seen, right / seen
print('%8s %10s %12s %12s' % ('batch', 'lr', 'val loss', 'val acc'))
for bs, lr in [(32, 0.05), (32, 0.1), (128, 0.1), (128, 0.4),
(512, 0.1), (512, 1.6)]:
loss, acc = run(lambda p, lr=lr: torch.optim.SGD(p, lr=lr),
epochs=3, batch_size=bs)
print('%8d %10.2f %12.4f %12.4f' % (bs, lr, loss, acc))
batch lr val loss val acc
32 0.05 0.2690 0.9230
32 0.10 0.2139 0.9335
128 0.10 0.3283 0.9065
128 0.40 0.2665 0.9200
512 0.10 0.4511 0.8775
512 1.60 1.8504 0.2805

Gradient accumulation

When the batch you want does not fit in memory, run several small batches, add their gradients together, and step once. The arithmetic is identical to one large batch, which is exactly the behaviour week 1 described as a bug when it happened by accident.

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
torch.manual_seed(0)
loader = DataLoader(train_set, batch_size=16, shuffle=False)
loss_fn = nn.CrossEntropyLoss()

def grad_after(batch_size, accumulate):
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 10))
loader = DataLoader(train_set, batch_size=batch_size, shuffle=False)
it = iter(loader)
model.zero_grad()
for _ in range(accumulate):
xb, yb = next(it)
(loss_fn(model(xb), yb) / accumulate).backward()
return torch.cat([p.grad.flatten() for p in model.parameters()])

big = grad_after(64, 1)
accumulated = grad_after(16, 4)
print('one batch of 64 norm %.6f' % big.norm().item())
print('four batches of 16 norm %.6f' % accumulated.norm().item())
print('largest difference %.2e'
% (big - accumulated).abs().max().item())
one batch of 64 norm 4.976234
four batches of 16 norm 4.976234
largest difference 7.45e-08

Divide by the number of accumulation steps

CrossEntropyLoss already averages over its batch, so adding four batch means gives you four times the gradient you wanted. Dividing each loss by the accumulation count restores the arithmetic. Miss it and your effective learning rate is four times what you think it is, which usually shows up as a run that diverges only when you change the accumulation setting.

Gradient clipping

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))
loss_fn = nn.CrossEntropyLoss()
xb, yb = next(iter(DataLoader(train_set, batch_size=64, shuffle=True)))

model.zero_grad()
(loss_fn(model(xb), yb) * 500).backward() # a pathological gradient
before = torch.cat([p.grad.flatten() for p in model.parameters()]).norm()
clipped = torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)
after = torch.cat([p.grad.flatten() for p in model.parameters()]).norm()

print('norm before clipping %.4f' % before.item())
print('reported by clip_grad_norm_ %.4f' % clipped.item())
print('norm after clipping %.4f' % after.item())
print('\ndirection is preserved, only the length changes.')
norm before clipping 1058.4519
reported by clip_grad_norm_ 1058.4520
norm after clipping 1.0000

direction is preserved, only the length changes.

Clipping is close to mandatory for recurrent networks in week 8 and for transformers in week 10, where a single unlucky batch can produce a gradient thousands of times the usual size and destroy a run that was going well. A max norm of 1.0 is the conventional starting point.

Day 3 takeaway

Batch size controls gradient noise, so it and the learning rate have to be tuned together. Accumulate gradients when memory is the constraint, dividing the loss by the number of accumulation steps. Clip the gradient norm for anything recurrent or attention based.
Week 03 · Day 4 of 7

When Training Goes Wrong

The failure shapes, and measuring gradients instead of guessing

By 699 words

Training goes wrong in a small number of recognisable ways. Learning the shapes of the failures is worth more than any amount of hyperparameter search.

Loss becomes nan

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
torch.manual_seed(0)
model = nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))
opt = torch.optim.SGD(model.parameters(), lr=8.0) # far too high
loss_fn = nn.CrossEntropyLoss()

for step, (xb, yb) in enumerate(DataLoader(train_set, batch_size=64,
shuffle=True)):
opt.zero_grad()
loss = loss_fn(model(xb), yb)
loss.backward()
opt.step()
if step < 8 or loss.item() != loss.item():
print('step %d loss %s' % (step, loss.item()))
if loss.item() != loss.item():
break
step 0 loss 2.285597085952759
step 1 loss 50.645973205566406
step 2 loss 5039.51953125
step 3 loss 99568.6171875
step 4 loss 3233553.75
step 5 loss 50211.375
step 6 loss 431.63970947265625
step 7 loss 1273.14404296875
SymptomMost likely causeFirst thing to try
Loss becomes nan or infLearning rate too highDivide the rate by ten
Loss is flat from step zeroRate far too small, or the graph is brokenOverfit one batch; check the loss is connected to the parameters
Loss falls then plateaus highModel too small, or the rate has decayed too farMore capacity, or check the schedule
Training loss falls, validation risesOverfittingWeek 4, in its entirety
Validation is wildly noisy between epochsValidation set too small, or the rate is too high late onLarger split, decay the rate
Accuracy is exactly the majority class rateThe model has collapsed to one outputCheck class balance and the loss weighting

Watching the gradients rather than guessing

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
torch.manual_seed(0)

def gradient_report(depth, activation):
layers = [nn.Flatten(), nn.Linear(784, 64), activation()]
for _ in range(depth):
layers += [nn.Linear(64, 64), activation()]
layers += [nn.Linear(64, 10)]
model = nn.Sequential(*layers)
xb, yb = next(iter(DataLoader(train_set, batch_size=64, shuffle=True)))
nn.CrossEntropyLoss()(model(xb), yb).backward()
norms = [p.grad.norm().item() for name, p in model.named_parameters()
if 'weight' in name]
return norms

for label, act in [('sigmoid', nn.Sigmoid), ('relu', nn.ReLU)]:
norms = gradient_report(8, act)
print('%-8s first layer %.2e last layer %.2e ratio %.1e'
% (label, norms[0], norms[-1], norms[-1] / max(norms[0], 1e-30)))
sigmoid first layer 8.89e-08 last layer 6.04e-01 ratio 6.8e+06
relu first layer 1.03e-03 last layer 7.58e-02 ratio 7.3e+01

This is the vanishing gradient, measured

With sigmoid activations the gradient reaching the first layer is orders of magnitude smaller than the one at the last, so the early layers barely move while the late ones train. With ReLU the ratio is far closer to one. Printing gradient norms per layer takes four lines and tells you immediately whether a deep model is training everywhere or only at the end.

A checklist that finds most problems

  1. Can the model overfit a single batch to near-zero loss? If not, the bug is structural, not statistical.
  2. Is the loss connected to the parameters? Print loss.grad_fn and check a parameter's .grad is not None after backward.
  3. Are the inputs normalised, and did you check the actual min and max rather than assuming?
  4. Are labels in the range the loss expects? CrossEntropyLoss wants 0 to C-1, and an off-by-one gives an unhelpful index error or, worse, silently trains on the wrong classes.
  5. Is model.train() set for training and model.eval() for evaluation?
  6. Did you call zero_grad()?
  7. Print the gradient norm per layer once, early.
  8. Try one tenth of the learning rate before you try anything clever.

Day 4 takeaway

A nan loss means the rate is too high, almost always. A flat loss means the model is disconnected or the rate is far too low. Measure gradient norms per layer rather than guessing at vanishing gradients, and work through the checklist before changing the architecture.
Week 03 · Day 5 of 7

Precision and Weight Averaging

Two cheap techniques worth knowing before you need them

By 805 words

Two techniques that cost almost nothing and are worth knowing before you need them: making training faster with lower precision, and making a model better by averaging the weights you already have.

Precision

import torch

x = torch.tensor([1.0], dtype=torch.float32)
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
info = torch.finfo(dtype)
print('%-16s bits %2d max %.3e smallest normal %.3e'
% (str(dtype), info.bits, info.max, info.tiny))

print('\na gradient of 1e-8 in each format:')
small = 1e-8
for dtype in [torch.float32, torch.float16, torch.bfloat16]:
print(' %-16s %s' % (str(dtype),
torch.tensor(small, dtype=dtype).item()))
torch.float32 bits 32 max 3.403e+38 smallest normal 1.175e-38
torch.float16 bits 16 max 6.550e+04 smallest normal 6.104e-05
torch.bfloat16 bits 16 max 3.390e+38 smallest normal 1.175e-38

a gradient of 1e-8 in each format:
torch.float32 9.99999993922529e-09
torch.float16 0.0
torch.bfloat16 1.0011717677116394e-08

Half precision has a much smaller range, and gradients are exactly where small numbers live. A gradient that underflows to zero in float16 contributes nothing, which is why mixed precision training scales the loss up before the backward pass and scales the gradients back down afterwards. On a GPU this roughly halves memory and can double speed; on the CPU this course runs on, it does neither, so the code below is shown rather than raced.

# The standard mixed precision loop on a CUDA device.
scaler = torch.amp.GradScaler('cuda')

for xb, yb in loader:
opt.zero_grad()
with torch.autocast('cuda', dtype=torch.float16):
loss = loss_fn(model(xb), yb)
scaler.scale(loss).backward() # scale up, so gradients survive
scaler.unscale_(opt) # scale back down before clipping
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
scaler.step(opt)
scaler.update()

bfloat16 is usually the better choice now

It has the same exponent range as float32 and fewer mantissa bits, so it does not underflow and needs no loss scaling at all. On hardware that supports it, torch.autocast('cuda', dtype=torch.bfloat16) with no GradScaler is simpler and less fragile. float16 remains common only because older cards support it and bfloat16 came later.

Averaging weights

Stochastic weight averaging: Train as usual, then average the weights from the last several epochs rather than taking the final ones. The average tends to sit in a flatter part of the loss surface, which generalises slightly better. It costs one extra copy of the model and no extra training.
import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
def build(seed=0):
torch.manual_seed(seed)
return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))

def run(opt_fn, epochs=3, batch_size=64, seed=0, sched_fn=None):
torch.manual_seed(seed)
model = build(seed)
opt = opt_fn(model.parameters())
sched = sched_fn(opt) if sched_fn else None
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val = DataLoader(val_set, batch_size=512)
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
if sched is not None:
sched.step()
model.eval()
right = seen = 0
total = 0.0
with torch.no_grad():
for xb, yb in val:
out = model(xb)
total += loss_fn(out, yb).item() * yb.numel()
right += (out.argmax(1) == yb).sum().item()
seen += yb.numel()
return total / seen, right / seen
import copy
torch.manual_seed(0)

model = build(0)
opt = torch.optim.SGD(model.parameters(), lr=0.05, momentum=0.9)
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=64, shuffle=True)
val = DataLoader(val_set, batch_size=512)

def score(m):
m.eval()
right = seen = 0
with torch.no_grad():
for xb, yb in val:
right += (m(xb).argmax(1) == yb).sum().item()
seen += yb.numel()
return right / seen

averaged = None
collected = 0
for epoch in range(1, 9):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
if epoch >= 5: # average the last four
state = model.state_dict()
if averaged is None:
averaged = {k: v.clone().float() for k, v in state.items()}
else:
for k in averaged:
averaged[k] += state[k].float()
collected += 1

final_acc = score(model)
for k in averaged:
averaged[k] /= collected
swa = copy.deepcopy(model)
swa.load_state_dict(averaged)
print('final epoch weights %.4f' % final_acc)
print('average of last four %.4f' % score(swa))
final epoch weights 0.9505
average of last four 0.9510

The averaged weights come out very slightly ahead here, by far less than the seed-to-seed variation day 6 measures on this same model. So the honest reading is that it did not help on this problem, and the reason to know about it is that it costs one extra copy of the weights and does reliably help on larger models trained for longer, where the final weights are still bouncing around a minimum rather than sitting in one.

Day 5 takeaway

Mixed precision halves memory and speeds up training on a GPU; prefer bfloat16 where it exists, because it needs no loss scaling. Weight averaging is close to free, and on a small model like this one it buys nothing you could distinguish from noise.
Week 03 · Day 6 of 7

Choosing Hyperparameters

What to search, on what budget, and the variation you must beat

By 852 words

Choosing hyperparameters for a deep model is the same problem as week 14 of the Machine Learning course, with one difference that changes the arithmetic completely: a single evaluation costs minutes or hours rather than milliseconds.

What is worth searching

ParameterPrioritySensible range
Learning rateFirst, and by a distanceFound by a range test, then tuned within a factor of three
Batch sizeSecond, jointly with the rateAs large as memory allows, then tune the rate to match
Weight decayThird1e-5 to 1e-1, on a log scale
Architecture width and depthFourthPowers of two; go wider before deeper
Dropout rateFifth0.0 to 0.5
Optimiser choiceRarely worth itAdamW unless you have a reason
Adam's betas and epsilonAlmost neverLeave them

Random search, on a budget

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
def build(seed=0):
torch.manual_seed(seed)
return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))

def run(opt_fn, epochs=3, batch_size=64, seed=0, sched_fn=None):
torch.manual_seed(seed)
model = build(seed)
opt = opt_fn(model.parameters())
sched = sched_fn(opt) if sched_fn else None
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val = DataLoader(val_set, batch_size=512)
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
if sched is not None:
sched.step()
model.eval()
right = seen = 0
total = 0.0
with torch.no_grad():
for xb, yb in val:
out = model(xb)
total += loss_fn(out, yb).item() * yb.numel()
right += (out.argmax(1) == yb).sum().item()
seen += yb.numel()
return total / seen, right / seen
import math
torch.manual_seed(0)
rng = torch.Generator().manual_seed(7)

def sample():
log_lr = -4 + 3 * torch.rand(1, generator=rng).item() # 1e-4 to 1e-1
wd = 10 ** (-5 + 4 * torch.rand(1, generator=rng).item())
return 10 ** log_lr, wd

print('%10s %10s %12s %12s' % ('lr', 'weight decay', 'val loss', 'val acc'))
results = []
for _ in range(8):
lr, wd = sample()
loss, acc = run(lambda p, lr=lr, wd=wd:
torch.optim.AdamW(p, lr=lr, weight_decay=wd),
epochs=2)
results.append((acc, lr, wd))
print('%10.5f %10.5f %12.4f %12.4f' % (lr, wd, loss, acc))

best = max(results)
print('\nbest: lr %.5f, weight decay %.5f, accuracy %.4f'
% (best[1], best[2], best[0]))
lr weight decay val loss val acc
0.00403 0.00006 0.2695 0.9220
0.00950 0.00424 0.3458 0.9070
0.00050 0.00050 0.3538 0.8955
0.00042 0.00330 0.3647 0.8940
0.00125 0.02541 0.3218 0.9035
0.03671 0.00160 0.6686 0.8360
0.00073 0.00007 0.3340 0.9045
0.00216 0.00027 0.2960 0.9145

best: lr 0.00403, weight decay 0.00006, accuracy 0.9220

Search on a short budget, then verify at full length

Two epochs is not enough to decide which model is best, but it is usually enough to decide which learning rates are hopeless, and that is most of the search space. Screen widely and cheaply, take the best three or four settings, and train those properly. The failure mode to watch for is a setting that is slow to start and best in the end, which a short screen discards; a warmup and a fixed number of steps rather than epochs both reduce it.

The number you report

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
def build(seed=0):
torch.manual_seed(seed)
return nn.Sequential(nn.Flatten(), nn.Linear(784, 128), nn.ReLU(),
nn.Linear(128, 10))

def run(opt_fn, epochs=3, batch_size=64, seed=0, sched_fn=None):
torch.manual_seed(seed)
model = build(seed)
opt = opt_fn(model.parameters())
sched = sched_fn(opt) if sched_fn else None
loss_fn = nn.CrossEntropyLoss()
loader = DataLoader(train_set, batch_size=batch_size, shuffle=True)
val = DataLoader(val_set, batch_size=512)
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
opt.step()
if sched is not None:
sched.step()
model.eval()
right = seen = 0
total = 0.0
with torch.no_grad():
for xb, yb in val:
out = model(xb)
total += loss_fn(out, yb).item() * yb.numel()
right += (out.argmax(1) == yb).sum().item()
seen += yb.numel()
return total / seen, right / seen
torch.manual_seed(0)
print('the same configuration, five different seeds:')
accs = []
for seed in range(5):
loss, acc = run(lambda p: torch.optim.SGD(p, lr=0.1, momentum=0.9),
epochs=2, seed=seed)
accs.append(acc)
print(' seed %d %.4f' % (seed, acc))
t = torch.tensor(accs)
print('\nmean %.4f, std %.4f, range %.4f'
% (t.mean(), t.std(), t.max() - t.min()))
print('any comparison smaller than that range is not a result.')
the same configuration, five different seeds:
seed 0 0.8995
seed 1 0.9000
seed 2 0.8910
seed 3 0.8825
seed 4 0.8915

mean 0.8929, std 0.0072, range 0.0175
any comparison smaller than that range is not a result.

Day 6 takeaway

Tune the learning rate first and almost everything else barely at all. Screen cheaply on short runs, then verify the survivors properly. And measure your seed-to-seed variation once, because it sets the smallest difference you are entitled to call an improvement.
Week 03 · Day 7 of 7

A Tuned Training Recipe

Every change measured on its own, against the week 2 baseline

By 906 words

Everything from this week applied at once, and then measured against the plain baseline from week 2 so the total is honest.

The tuned run

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
import time
torch.manual_seed(0)

EPOCHS, BATCH = 8, 128
loader = DataLoader(train_set, batch_size=BATCH, shuffle=True)
val = DataLoader(val_set, batch_size=512)
steps = EPOCHS * len(loader)

model = nn.Sequential(nn.Flatten(), nn.Linear(784, 256), nn.ReLU(),
nn.Linear(256, 128), nn.ReLU(), nn.Linear(128, 10))
opt = torch.optim.AdamW(model.parameters(), lr=3e-3, weight_decay=1e-4)
sched = torch.optim.lr_scheduler.OneCycleLR(opt, max_lr=3e-3,
total_steps=steps)
loss_fn = nn.CrossEntropyLoss()

def score():
model.eval()
right = seen = 0
with torch.no_grad():
for xb, yb in val:
right += (model(xb).argmax(1) == yb).sum().item()
seen += yb.numel()
return right / seen

start = time.time()
for epoch in range(1, EPOCHS + 1):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
opt.step()
sched.step()
if epoch % 2 == 0:
print('epoch %d lr %.5f val acc %.4f'
% (epoch, sched.get_last_lr()[0], score()))
print('\n%.1f seconds' % (time.time() - start))
epoch 2 lr 0.00282 val acc 0.8970
epoch 4 lr 0.00242 val acc 0.9270
epoch 6 lr 0.00084 val acc 0.9495
epoch 8 lr 0.00000 val acc 0.9505

11.5 seconds

Against the baseline

import torch
from torch import nn
from torch.utils.data import DataLoader, Subset
from torchvision import datasets, transforms

tf = transforms.Compose([transforms.ToTensor(),
transforms.Normalize((0.1307,), (0.3081,))])
train_all = datasets.MNIST('data', train=True, download=True, transform=tf)
train_set = Subset(train_all, range(8000))
val_set = Subset(train_all, range(8000, 10000))
import time

def train(tag, opt_fn, epochs, batch, sched_fn=None, clip=None,
hidden=(128,)):
torch.manual_seed(0)
sizes = (784,) + hidden
layers = [nn.Flatten()]
for a, b in zip(sizes, sizes[1:]):
layers += [nn.Linear(a, b), nn.ReLU()]
layers += [nn.Linear(sizes[-1], 10)]
model = nn.Sequential(*layers)
loader = DataLoader(train_set, batch_size=batch, shuffle=True)
val = DataLoader(val_set, batch_size=512)
opt = opt_fn(model.parameters())
sched = sched_fn(opt, epochs * len(loader)) if sched_fn else None
loss_fn = nn.CrossEntropyLoss()
start = time.time()
for _ in range(epochs):
model.train()
for xb, yb in loader:
opt.zero_grad()
loss_fn(model(xb), yb).backward()
if clip:
torch.nn.utils.clip_grad_norm_(model.parameters(), clip)
opt.step()
if sched:
sched.step()
model.eval()
right = seen = 0
with torch.no_grad():
for xb, yb in val:
right += (model(xb).argmax(1) == yb).sum().item()
seen += yb.numel()
print('%-30s %10.4f %10.1fs' % (tag, right / seen, time.time() - start))

print('%-30s %10s %11s' % ('', 'val acc', 'time'))
train('week 2 baseline, SGD 0.1', lambda p: torch.optim.SGD(p, lr=0.1),
8, 64)
train('+ momentum',
lambda p: torch.optim.SGD(p, lr=0.1, momentum=0.9), 8, 64)
train('+ AdamW',
lambda p: torch.optim.AdamW(p, lr=3e-3, weight_decay=1e-4), 8, 64)
train('+ one cycle and clipping',
lambda p: torch.optim.AdamW(p, lr=3e-3, weight_decay=1e-4), 8, 128,
sched_fn=lambda o, s: torch.optim.lr_scheduler.OneCycleLR(
o, max_lr=3e-3, total_steps=s), clip=1.0)
train('+ a wider model',
lambda p: torch.optim.AdamW(p, lr=3e-3, weight_decay=1e-4), 8, 128,
sched_fn=lambda o, s: torch.optim.lr_scheduler.OneCycleLR(
o, max_lr=3e-3, total_steps=s), clip=1.0, hidden=(256, 128))
val acc time
week 2 baseline, SGD 0.1 0.9375 10.1s
+ momentum 0.9200 10.1s
+ AdamW 0.9425 10.5s
+ one cycle and clipping 0.9410 10.5s
+ a wider model 0.9495 10.4s

Now compare that spread against day 6

Five seeds of one configuration varied by 0.0175 in accuracy. The entire table above spans about that much, and the row labelled + momentum is below the baseline it was supposed to improve on. So the correct conclusion from this experiment is not “this recipe is worth three points”. It is that on eight thousand MNIST digits with a small dense network, none of these changes reliably matters except making the model wider.

That is a real finding and it generalises: optimiser tuning pays when the model is large, the dataset is hard and the run is long. On an easy problem it mostly buys you the ability to be careless about the learning rate. Week 5 changes the architecture instead, and the difference there is not subtle.

Add one thing at a time and keep the table

This is the only way to know what your training recipe is actually made of. Applied all at once, five changes produce one number and no understanding, and when you later need to cut the training time in half you will not know which of them you can drop. Keep the table. It is also the honest answer when somebody asks why your model has a one cycle schedule in it.

The recipe worth starting from

  1. AdamW, learning rate 3e-3 for a small model or 1e-3 for a large one, weight decay 1e-4.
  2. Batch size as large as memory allows.
  3. One cycle or cosine schedule over the whole run.
  4. A few hundred steps of warmup if the model is large.
  5. Gradient clipping at norm 1.0.
  6. Then run the range test and adjust the rate, because that is the one that varies most between problems.

Day 7 takeaway

Optimiser, learning rate, schedule, batch size and clipping are one system rather than five choices. Start from a known recipe, change one thing at a time, keep the table of what each was worth, and measure your seed variation so you know which differences are real.