The Learning Rate
The most important dial.
Ask any experienced deep learning engineer which setting they'd tune if they could only choose one. They'll all give the same answer. It's the learning rate. It's also the first thing to check when training goes wrong.
- If you could tune just one
- Everyone gives the same answer
- First thing to check
Four learning rates
We train one epoch, with four different rates. Run it. Look at the results. Point zero zero zero zero one is so small that the model has barely learned. Point zero zero one works well. Point zero one is still fine, on this easy problem. And at one, training has collapsed completely, to a score close to random guessing. Far too big, and the weights fly apart. That's what a broken learning rate looks like.
import torch
from torch import nn
from torchvision import datasets, transforms
root = "data"
tf = transforms.ToTensor()
train = datasets.FashionMNIST(root, train=True, transform=tf)
test = datasets.FashionMNIST(root, train=False, transform=tf)
X_test = test.data.float().unsqueeze(1) / 255
def run(make_opt, epochs=2, seed=0, net=None, scheduler=None):
torch.manual_seed(seed)
model = net() if net else nn.Sequential(nn.Flatten(), nn.Linear(784, 256),
nn.ReLU(), nn.Linear(256, 10))
opt = make_opt(model.parameters())
dl = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True)
for _ in range(epochs):
model.train()
for x, y in dl:
loss = nn.functional.cross_entropy(model(x), y)
opt.zero_grad(); loss.backward(); opt.step()
model.eval()
with torch.no_grad():
return (model(X_test).argmax(1) == test.targets).float().mean().item()
from torch.optim import Adam
for lr in [1e-5, 1e-3, 1e-2, 1.0]:
print(f"lr {lr:<7} accuracy {run(lambda p: Adam(p, lr=lr), epochs=1):.3f}")lr 1e-05 accuracy 0.652 lr 0.001 accuracy 0.842 lr 0.01 accuracy 0.818 lr 1.0 accuracy 0.181
Warm up, then cool down
The best learning rate actually changes over time. At the start, the weights are random, so big jumps can be dangerous. We warm up slowly, for a short time. Then, as we approach a good solution, we need finer and finer steps. So we slowly cool down, often along a smooth cosine curve. This warm up and cosine schedule is used to train nearly all modern large models.
A schedule in code
PyTorch has schedulers built in. We create one for the linear warm up, and one for the cosine decay, and chain them together. Each training step, we call sched dot step. Run it. The rate climbs for the first ten steps, then glides down towards zero. You'd normally use this inside the training loop, right after the optimiser step.
import torch
from torch import nn
model = nn.Linear(10, 1)
opt = torch.optim.AdamW(model.parameters(), lr=0.001)
steps = 100
warm = torch.optim.lr_scheduler.LinearLR(opt, 0.1, 1.0, total_iters=10)
cos = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=steps - 10)
sched = torch.optim.lr_scheduler.SequentialLR(opt, [warm, cos], milestones=[10])
for step in range(steps):
opt.step()
sched.step()
if step in (0, 9, 30, 60, 99):
print(f"step {step:>2} lr {opt.param_groups[0]['lr']:.6f}")step 0 lr 0.000190 step 9 lr 0.001000 step 30 lr 0.000872 step 60 lr 0.000396 step 99 lr 0.000000
Summary
Too small, and you wait forever. Too large, and everything explodes. And a warm up with a cosine decay gets the best of both. Next, a hidden trap that can ruin training before the first step. How you start your weights.
Key takeaways- Tune the learning rate first
- Too big explodes, too small crawls
- Warm up, then cosine decay