Overfit One Batch

You start a training run on Friday evening. On Monday morning, the accuracy is still rubbish. Was it a bad learning rate? A bug in the data? Is the model too small? You've wasted a whole weekend, and you still don't know. There's a test which takes one minute, and it could have saved you.

• Training all weekend
• The accuracy is rubbish
• A test that takes one minute

Memorise thirty two images

Take just thirty two images, and train on them over and over. Any working model must be able to memorise thirty two images. This isn't about being good. It's a test of the plumbing. Run it. The loss goes from about two point three, which is what random guessing gives you with ten classes, down to almost zero. So the data, the labels, the model, the loss, and the optimiser are all connected properly. If it hadn't reached zero, there'd be a bug. And you'd know that in sixty seconds, rather than on Monday.

import torch
from base import load, net

torch.manual_seed(0)
X, y, _, _ = load(32)
model = net()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
for step in range(301):
    loss = torch.nn.functional.cross_entropy(model(X), y)
    opt.zero_grad(); loss.backward(); opt.step()
    if step % 100 == 0:
        print(f"Step {step}: loss {loss.item():.4f}")

Step 0: loss 2.2920
Step 100: loss 0.0067
Step 200: loss 0.0023
Step 300: loss 0.0012

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()

What the failure tells you

If this test fails, the way that it fails is a clue. If the loss doesn't move at all, gradients aren't flowing. Check for a detach, or a forgotten optimiser step. If the loss explodes, the learning rate is too high. If it settles at a high value, then your labels may be misaligned with your inputs. And if it starts far above two point three, your outputs or your data scaling are off.

Summary