The Training Loop and Checkpoints

This section puts the pieces together: a training loop in spotter.ss, checkpoints to save and resume it, and a short Python script that only provides the pictures. Then it makes the loop work for more than one model, and swaps it for the standard library’s.

The program

First, one step of training that also reports. It replaces Section 2’s step:

## One step of training.
fn train_step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> StepLogs:
    l, g = value_and_grad(loss, m, batch, mode)
    n = norm(g)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return StepLogs(loss=l, accuracy=accuracy(m, batch), grad_norm=n)

Delete show_accuracy and log_step from Section 2: train_step and the loop do their jobs. Then the loop, with checkpoints:

## Trains a new spotter for `epochs` epochs.
fn fit(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    opt = adamw(m, lr=3e-3)
    run(mut m, mut opt, train, sink, 0, epochs, mode)

## Carries on from the checkpoint of the last epoch that finished.
fn resume(train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode) -> Spotter:
    m, opt, done = load("ckpt.pt", Spotter, AdamW[Spotter], int)
    run(mut m, mut opt, train, sink, done, epochs, mode)
    return m

## Epochs `first` to `epochs`, with a checkpoint after each.
fn run(mut m: Spotter, mut opt: AdamW[Spotter], train: Dataset[Shape], sink: LogSink,
                first: int, epochs: int, mode: Mode):
    for epoch in first..epochs:
        for batch in train |> shuffle(seed=epoch) |> batches(64):
            sink.write(train_step(mut m, mut opt, batch, mode))
        sink.flush()
        save("ckpt.pt", m, opt, epoch + 1)
        print("epoch", epoch, "done")

Read run from the top: each epoch shuffles with its own seed, writes a log record per step, flushes the sink (one CSV row per epoch), and saves a checkpoint. Nothing happens that the code doesn’t say. That’s the whole of Chapter 10’s Python loop, the optimizer included, now checked: train_step can’t be given an optimizer for another model, a batch of the wrong record, or a misspelled metric.

Checkpoints

save(path, values...) writes models, optimizer state (moments, learning rate, step count) and host numbers, such as the number of epochs done. load(path, Types...) takes the type of each value in the file and checks the file against them before it makes anything: a Spotter saved before you changed D is an error naming each field whose shape differs, not a w[0] as [P*P D] that fails on the first call. Loaded tensors go on the program’s device, never the one they were saved from.

Only models, optimizer state and numbers can be saved. A tensor that’s part of the training state belongs in a model, as a buffer if it isn’t trained:

dim S = 16

fn keep(img: [b 1 S S]):
    save("img.pt", img)

Running it

The Python side is now just the pictures. Shapes wraps Chapter 7’s draw as a dataset, with each item a dict of the record’s fields:

import shapescript
from shapes import draw
from spotter import Spotter, fit, resume, predict


class Shapes:
    """`n` labelled pictures from `draw`, for ShapeScript's loader."""

    def __init__(self, n, seed):
        self.img, self.label = draw(n, seed)

    def __len__(self):
        return len(self.label)

    def __getitem__(self, i):
        return {"img": self.img[i], "label": int(self.label[i])}


shapescript.loader_workers(0)  # the dataset is defined in this script
train = Shapes(2000, seed=0)
sink = shapescript.sinks.csv("log.csv", every=None)
model = Spotter.init(seed=0)
fit(model, train=train, sink=sink, epochs=4, mode=shapescript.Train(seed=0))
model = resume(train=train, sink=sink, epochs=8, mode=shapescript.Train(seed=0))
sink.close()

test = Shapes(500, seed=1)
right = (predict(model, test.img) == test.label).float().mean().item()
print(f"test accuracy: {right:.1%}")

$ shapescript run train.py
epoch 0 done
epoch 1 done
epoch 2 done
epoch 3 done
epoch 4 done
epoch 5 done
epoch 6 done
epoch 7 done
test accuracy: 97.2%
$ cat log.csv
loss,accuracy,grad_norm
1.0460098292678595,0.59521484375,1.4748300313949585
0.501040487550199,0.8466796875,2.4373481273651123
0.26240775594487786,0.91845703125,2.432314395904541
0.1644725656369701,0.962890625,2.5464417934417725
0.08703674498246983,0.984375,2.1427037715911865
0.040361865219892934,0.99560546875,1.6711437702178955
0.028169181416160427,0.99853515625,1.5459754467010498
0.010340657085180283,1.0,1.0254361629486084

fit trained four epochs and saved after each; resume loaded the fourth checkpoint and carried on to eight. The result is exactly that of training eight epochs without stopping, dropout included: each dropout mask is drawn from the mode, the epoch and the batch, not from how many steps this run has taken, so a run that starts at epoch 4 draws the masks of epoch 4. fit(model, ...) changed model in place: from Python, a mut argument is changed, as a PyTorch module would be. loader_workers(0) loads batches in this process, because a dataset class defined in the script you run can’t be sent to worker processes.

Read the log’s last row against the test accuracy. In the last epoch, the spotter got every training picture right, but nearly 3% of new ones wrong. It has started to learn its 2,000 pictures by heart, which is what dropout is there to slow down; more pictures, or the nudged ones of Chapter 10, would help more.

One loop for every model

train_step only trains a Spotter. Is the spotter worth its mixer blocks? A fair test trains a much simpler model on the same batches. Two features make one loop work for both.

Overloading by model. Several functions may share a name if each takes a different model first. So each model brings its own forward and loss, and a call picks the one for its model. In a new file, compare.ss, next to spotter.ss:

import spotter

## A baseline: one matrix from pixels to scores.
model Linear:
    param w: [spotter.S*spotter.S spotter.C] = zeros

fn loss(m: Linear, batch: [b] spotter.Shape, mode: Mode) -> [] NonNeg:
    x = batch.img |> to Unit |> rearrange(b c h w -> b (c h w))
    return cross_entropy(x @ m.w, batch.label, over=spotter.C)

loss now has two versions, spotter.loss for a Spotter and this one for a Linear.

Generic functions. A parameter type that’s a single capital letter, such as M or E, that nothing declares, is a type variable: m: M takes any model, and [b] E a batch of any data record. A function with one is checked again at each call, for the types it’s given, so it needs only what it uses. There are no traits or interfaces to declare:

## One SGD step on any model that has a `loss`.
fn step(mut m: M, mut opt: SGD[M], batch: [b] E, mode: Mode) -> []:
    l, g = value_and_grad(loss, m, batch, mode)
    update(mut m, mut opt, g)
    return l

## The losses of both models, trained side by side for an epoch.
fn race(mut a: spotter.Spotter, mut z: Linear, data: Dataset[spotter.Shape], mode: Mode):
    oa = sgd(a, lr=0.1)
    oz = sgd(z, lr=0.1)
    for batch in data |> shuffle(seed=0) |> batches(64):
        print(step(mut a, mut oa, batch, mode), step(mut z, mut oz, batch, mode))

At step(mut a, ...), M is spotter.Spotter and loss(m, ...) is the spotter’s; at step(mut z, ...), it’s Linear’s. A model without the function a generic one needs is an error at the call, naming what’s missing and where it’s used.

The standard library’s fit

The standard library has a generic fit of its own, written in ShapeScript like step above, with the batch size and learning rate as parameters. All it needs from a model is a loss(m, batch, mode), which the spotter has. Import it at the top of spotter.ss:

import shapescript.std.fit as std

and call it:

## Trains with the standard library's loop instead.
fn fit_std(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    std.fit(mut m, train, sink, epochs, 64, 3e-3, mode)

It’s short enough to read on one screen, and meant to be copied when you need something different: shapescript.std.path("fit") in Python gives the file. Read it, and you’ll find a train_step and a run, the names of the spotter’s own. They don’t clash, because each file is a namespace. Inside fit.ss, run means fit.ss’s run, whatever the model’s file declares. A generic function looks next to the model only for names its own file doesn’t declare, which is how std.fit finds the spotter’s loss. So you don’t need to know a library’s helpers to name your functions, and a library can add helpers without changing what your code means.

The spotter so far

The whole of spotter.ss, as the project leaves it:

devices cpu

import mixer
import shapescript.std.fit as std

## Side of each image, in pixels.
dim S = 16
## Side of each square patch.
dim P = 4
## Number of patches in an image.
dim N = (S / P) * (S / P)
## Width of each token.
dim D = 32
## Hidden width of the MLP that mixes patches.
dim Dt = 64
## Hidden width of the MLP that mixes channels.
dim Dc = 64
## Number of classes: square, disk, plus and cross.
dim C = 4
## Number of mixer blocks.
dim L = 2

## A batch of grayscale images.
dim Images = [b 1 S S]

## The spotter's weights.
model Spotter:
    param w_e: [P*P D]  = normal(std=0.25)
    param tw1: [L N Dt] = normal(std=0.25)
    param tw2: [L Dt N] = normal(std=0.125)
    param cw1: [L D Dc] = normal(std=0.18)
    param cw2: [L Dc D] = normal(std=0.125)
    param w_h: [D C]    = normal(std=0.18)

## One labelled picture.
data Shape:
    img:   [1 S S] u8
    label: [] Index[C]

## What each training step reports.
logs StepLogs:
    loss:      [] mean
    accuracy:  [] mean
    grad_norm: [] max

## The images that have any ink: a pixel brighter than 0.5.
fn nonblank(img: Images) -> [?k 1 S S]:
    return img[reduce(img, b c h w -> b, max) > 0.5]

## The average token of each image.
fn pool(x: [b n D]) -> [b D] = mean(x, over=n)

## Cuts each image into P×P patches, one row of pixels per patch.
fn patchify(img: Images) -> [b N P*P]:
    return img |> rearrange(b c (gh p1) (gw p2) -> b (gh gw) (p1 p2 c), p1=P, p2=P)

## A token per patch.
fn embed(img: Images, w_e: [P*P D]) -> [b N D] = patchify(img) @ w_e

## The class scores of each image.
fn forward(m: Spotter, img: Images u8, mode: Mode) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_tokens(m.tw1[l], m.tw2[l]) |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> pool |> dropout(0.1, mode) |> matmul(m.w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, batch: [b] Shape, mode: Mode) -> [] NonNeg:
    return cross_entropy(forward(m, batch.img, mode), batch.label, over=C)

## The likeliest class of each image.
fn predict(m: Spotter, img: Images u8) -> [b] Index[C]:
    return argmax(forward(m, img, Eval), over=C)

## The fraction of a batch predicted right.
fn accuracy(m: Spotter, batch: [b] Shape) -> []:
    return mean((predict(m, batch.img) == batch.label) to f32)

## The fraction of each class's images that were predicted right.
fn per_class_accuracy(pred: [b] Index[C], y: [b] Index[C]) -> [C]:
    return C |> map(c => (sum(pred == c and y == c) to f32) / maximum(sum(y == c) to f32, 1.0))

## How often each true class (rows) was predicted as each class (columns).
fn confusion(pred: [b] Index[C], y: [b] Index[C]) -> [C C] i64:
    return C |> map(t => C |> map(p => sum(y == t and pred == p)))

## Clears the head, to retrain it on new classes.
fn reset_head(mut m: Spotter):
    m.w_h = m.w_h * 0.0

## One step of training.
fn train_step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> StepLogs:
    l, g = value_and_grad(loss, m, batch, mode)
    n = norm(g)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return StepLogs(loss=l, accuracy=accuracy(m, batch), grad_norm=n)

## Trains a new spotter for `epochs` epochs.
fn fit(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    opt = adamw(m, lr=3e-3)
    run(mut m, mut opt, train, sink, 0, epochs, mode)

## Carries on from the checkpoint of the last epoch that finished.
fn resume(train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode) -> Spotter:
    m, opt, done = load("ckpt.pt", Spotter, AdamW[Spotter], int)
    run(mut m, mut opt, train, sink, done, epochs, mode)
    return m

## Epochs `first` to `epochs`, with a checkpoint after each.
fn run(mut m: Spotter, mut opt: AdamW[Spotter], train: Dataset[Shape], sink: LogSink,
                first: int, epochs: int, mode: Mode):
    for epoch in first..epochs:
        for batch in train |> shuffle(seed=epoch) |> batches(64):
            sink.write(train_step(mut m, mut opt, batch, mode))
        sink.flush()
        save("ckpt.pt", m, opt, epoch + 1)
        print("epoch", epoch, "done")

## Trains with the standard library's loop instead.
fn fit_std(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode):
    std.fit(mut m, train, sink, epochs, 64, 3e-3, mode)

## Rolls the axis `dims` round by `shifts` places.
extern fn roll(input: [...] u8, shifts: Size[K], dims: Axis[n]) -> [...] u8 = "torch.roll"

## Moves each image one pixel to the right; the last column wraps round.
fn nudge(img: [b c h w] u8) -> [b c h w] u8 = roll(img, 1, dims=w)

## The largest value of each `kernel_size`-wide square, every `stride` pixels.
extern fn max_pool2d(input: [b c h w], kernel_size: Size[K], stride: Size[St]) -> [b c (h - K)/St + 1 (w - K)/St + 1] = "torch.nn.functional.max_pool2d"

## A 7×7 summary of each image.
fn thumbnail(img: Images) -> [b 1 7 7] = max_pool2d(img, 3, stride=2)

From six loose tensors and a Python loop to a model that trains itself: that’s the spotter done. Chapters 13 to 16 are about the tools around a program like it, and Chapter 17 builds a second project, a small transformer, with everything at once.

Exercises

1. Make run print "saved" after every checkpoint. Where does the line go, and why must it be indented like save?
2. Write retrain_head(mut m: Spotter, train: Dataset[Shape], sink: LogSink, epochs: int, mode: Mode), which clears the head with reset_head and trains again with fit. Which muts does it need?
3. Train compare.ss’s Linear baseline with std.fit, in a function fit_linear in compare.ss. What does std.fit need from Linear that it already has?