Data, Modes and Logs

The spotter has weights and can take a step. To train it properly it needs three more things: some randomness while training, a dataset to draw batches from, and a way to report how it’s going.

Modes and dropout

A spotter that sees the same 2,000 pictures again and again can learn those pictures by heart instead of learning shapes. Dropout is the usual guard: while training, it zeroes a random tenth of the values, so the head can’t rely on any one of them. When the spotter is used, it must be off.

In PyTorch, model.train() and model.eval() flip a flag that dropout reads, and forgetting eval() before evaluating is a classic bug. In ShapeScript there’s no flag. A function that does something random takes a Mode argument, which is Train(seed=s) or Eval. Change forward to take one, and put dropout between the pooling and the head’s matrix product; head can go, since forward now does its two steps with dropout between them:

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

predict passes Eval, so whoever calls it can’t get it wrong. With Eval, dropout returns its input; with Train(seed=s), it zeroes each value with probability 0.1 and divides the others by 0.9.

A mode isn’t a flag. A bool named training would be the PyTorch habit, and the checker refuses it:

dim D = 32
dim C = 4

fn scores(x: [b D], w: [D C], training: bool) -> [b C]:
    return x |> dropout(0.1, training) |> matmul(w)

$ shapescript check flag.ss
error[E1202]: expected a `Mode`, found `bool`
 --> flag.ss:5:30
  |
5 |     return x |> dropout(0.1, training) |> matmul(w)
  |                              ^^^^^^^^ `dropout`'s parameter `mode` is a `Mode`
  |
  = help: pass `Train(seed=…)` to train or `Eval` to evaluate
  = note: in stage 1 of this pipeline, whose input is `[b D]`

For more information about E1202, run `shapescript explain E1202`.

The difference is the seed. A mode carries the seed that random operations draw from, so training is reproducible with no global random state. Each call that passes a mode on folds in its own name, so the two dropouts of a model with two draw different masks from one seed. And randomness shows in the signatures: dropout needs a mode, so forward takes one, and so do loss and everything else that calls it, with each call passing one on. A function without a Mode parameter, like predict, either passes Eval or isn’t random.

Data records

Until now, a batch was two arguments, img and y, that had to agree on b. A data declaration says what one example holds, so the two travel together:

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

[b] Shape is a batch of b examples. Every field gets the batch’s axis first: batch.img is [b 1 S S] u8, an Images u8, and batch.label is [b] Index[C]. So loss can take one batch instead of two tensors, and pass the mode on:

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

and step changes to match. Delete grads, which was there to show a gradient, and train, which only ever saw one batch:

## One step of training on a batch, giving the loss before it.
fn step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> [] NonNeg:
    l, g = value_and_grad(loss, m, batch, mode)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return l

A misspelled field is an error, and so is a field whose shape uses a generic size: every example of a record has the same shapes, so they can be stacked into batches. (Examples of different lengths, such as sentences, are padded to a dim, with a Mask field saying which positions are real. Padding them for you is coming.)

Datasets and batches

Dataset[Shape] is a dataset of Shapes. It comes from Python: any map-style dataset, an object with __len__ and __getitem__ such as a torch.utils.data.Dataset, whose items have the record’s fields, as a dict or a tuple. Section 3 wraps draw in one. Each item is checked against the record as it’s loaded, so a picture of the wrong size, or a label of 4, fails at once, naming the field.

Datasets are shuffled and batched with two pipe stages, and a for loop runs over the batches. Here’s how well the spotter does on each batch of a dataset:

## Prints the accuracy of each batch, in a shuffled order.
fn show_accuracy(m: Spotter, data: Dataset[Shape], epoch: int):
    for batch in data |> shuffle(seed=epoch) |> batches(64):
        print("accuracy", accuracy(m, batch))

shuffle(seed=epoch) gives an order decided by the seed alone. The last batch holds what’s left over, so the batch size isn’t known statically: inside the loop, batch is a [b] Shape whose b is an unknown size. (When it can, the loader pads the last batch to the full size and masks the padding out, so every step has the same shapes; means and losses still count only the real examples.)

Batches are built by worker processes and copied to the declared device for you.

Logs and sinks

A logs declaration names the numbers a program reports, and how each combines when several are reported together, such as over an epoch:

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

Over an epoch, a row will have the mean loss, the mean accuracy, and the largest gradient norm, the one to watch for a step that went wrong. StepLogs(loss=l, accuracy=a, grad_norm=n) makes a record, and every field must be given, by name. A misspelled metric is a compile error, not a missing column you notice a week later:

fn report(l: [] NonNeg) -> StepLogs = StepLogs(loss=l, accuracy=0.0, grad_nrom=0.0)

Records are written to a sink, a LogSink, which comes from Python: any object with a write(dict) method. The shapescript.sinks module has sinks for CSV files and TensorBoard that combine records by each field’s rule: shapescript.sinks.csv("log.csv", every=None) writes one row each time the program calls sink.flush().

## Reports one step.
fn log_step(sink: LogSink, l: [] NonNeg, a: [], n: [] NonNeg):
    sink.write(StepLogs(loss=l, accuracy=a, grad_norm=n))

print(a, b, ...), as in show_accuracy, prints any values, tensors included, copying them to the CPU as needed.

The spotter so far

devices cpu

import mixer

## 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 on a batch, giving the loss before it.
fn step(mut m: Spotter, mut opt: AdamW[Spotter], batch: [b] Shape, mode: Mode) -> [] NonNeg:
    l, g = value_and_grad(loss, m, batch, mode)
    update(mut m, mut opt, g |> clip_norm(1.0))
    return l

## Prints the accuracy of each batch, in a shuffled order.
fn show_accuracy(m: Spotter, data: Dataset[Shape], epoch: int):
    for batch in data |> shuffle(seed=epoch) |> batches(64):
        print("accuracy", accuracy(m, batch))

## Reports one step.
fn log_step(sink: LogSink, l: [] NonNeg, a: [], n: [] NonNeg):
    sink.write(StepLogs(loss=l, accuracy=a, grad_norm=n))

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

Exercises

1. Pictures in colour would need a different record. Write data ColorShape with an RGB image, [3 S S] u8, and the label, and give the type of batch.img for batch: [b] ColorShape.
2. Write count_right(m: Spotter, batch: [b] Shape) -> [], the number of pictures of a batch the spotter gets right.
3. Add a field lr to StepLogs, combined by its last value with last, and write log_lr(sink: LogSink, l: [] NonNeg, lr: float) that reports it with the loss. Which rule would you give a learning rate, and why not mean?