The Learning Rate Finder
One minute, one good guess.
The learning rate is the most important setting that you've got. Most people guess. They try one value, wait for hours, and then try another. In twenty fifteen, Leslie Smith suggested something much smarter. Don't train several times. Train once, and raise the learning rate as you go.
- The most important setting
- Most people just guess
- Train once, and raise it as you go
Finding the edge
We start with an extremely small rate. After every batch, we multiply it by a tiny bit more. The loss falls, and at some point, the rate gets so large that the loss blows up. We stop right there. Run it. It tells us a sensible starting rate, of about point zero one. We've learned, in one pass through the data, roughly where the safe limit is. The common advice is to use a value about ten times smaller than the point where it blew up. That's what we've printed.
import torch
from base import load, net
torch.manual_seed(0)
X, y, _, _ = load()
model = net()
opt = torch.optim.SGD(model.parameters(), lr=1e-5)
lr, best = 1e-5, 9
for i in torch.randperm(len(X)).split(32):
loss = torch.nn.functional.cross_entropy(model(X[i]), y[i])
opt.zero_grad(); loss.backward(); opt.step()
if len(i) == 32 and i[0] % 1 == 0:
best = min(best, loss.item())
if loss.item() > 4 * best:
print(f"Loss blew up at learning rate {lr:.3f}")
break
lr *= 1.05
for g in opt.param_groups:
g["lr"] = lr
print(f"Good starting rate: about {lr / 10:.3f}")Good starting rate: about 0.010
Also uses base.py
import torch
from torchvision import datasets, transforms
def load(n=6000):
t = transforms.ToTensor()
a = datasets.MNIST("data", True, t, download=False)
b = datasets.MNIST("data", False, t, download=False)
return a.data[:n].float()[:, None] / 255, a.targets[:n], b.data[:2000].float()[:, None] / 255, b.targets[:2000]
def net():
return torch.nn.Sequential(torch.nn.Flatten(), torch.nn.Linear(784, 128), torch.nn.ReLU(),
torch.nn.Linear(128, 10))
def fit(model, X, y, lr=1e-3, epochs=3, opt=None):
opt = opt or torch.optim.Adam(model.parameters(), lr=lr)
for _ in range(epochs):
for i in torch.randperm(len(X)).split(64):
loss = torch.nn.functional.cross_entropy(model(X[i]), y[i])
opt.zero_grad(); loss.backward(); opt.step()
return model
def score(model, X, y):
return (model(X).argmax(1) == y).float().mean().item()Summary
One quick sweep beats hours of guessing. But what about all of the other settings? Next, a search strategy that surprises people.
Key takeaways- Raise the rate each batch
- Stop when the loss explodes
- Divide by ten