Logging Every Run

Three weeks ago, you got your best ever result. Which settings was it? You can't remember. There's a folder of files named test two, and new final. Every experienced practitioner has lost a great result this way. Just once.

• Your best ever result
• Which settings was it?
• Everyone loses one, once

Logging the signals that matter

Here's a simple training loop which writes to a file. We record the loss, and also the gradient norm, which measures how large the gradients are. Clip grad norm returns that number, and caps it too. Run it. A small, clean record for each epoch. The gradient norm is an underrated early warning. If it suddenly spikes, then the training's about to go wrong. If it falls to zero, gradients have died. This one number has saved a lot of training runs.

import csv, torch
from base import load, net

torch.manual_seed(0)
X, y, Xt, yt = load(3000)
model = net()
opt = torch.optim.Adam(model.parameters(), lr=1e-3)
with open("run.csv", "w", newline="") as f:
    log = csv.writer(f)
    log.writerow(["epoch", "loss", "grad_norm"])
    for epoch in range(3):
        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()
            norm = torch.nn.utils.clip_grad_norm_(model.parameters(), 5.0)
            opt.step()
        log.writerow([epoch, round(loss.item(), 3), round(norm.item(), 2)])
print(open("run.csv").read().strip())

epoch,loss,grad_norm
0,0.634,1.41
1,0.356,0.69
2,0.433,1.04

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 to record

For every run, save the code version, all settings, the random seed, and the data version. Log the training and validation loss, the learning rate, and the gradient norm. Keep a few sample predictions, too. Tools such as TensorBoard and Weights and Biases do this for you, and they're worth it after your first ten runs.

Summary