Models, mut and Gradients

This section turns the spotter’s six weight tensors into one value, a Spotter, and gives it the two things training needs: a gradient, and an optimizer to follow it.

Models

A model declares the tensors that make up a network’s state. In spotter.ss, below the dims, add:

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

• A param is a tensor that’s trained. Its type is a shape, as always, but one that uses only numbers and dims, so every Spotter has the same shapes.
• Every field has an initializer: zeros, ones, full(value), normal(mean=0, std=1) or uniform(lo=0, hi=1). The standard deviations here are about 1 / √(the axis each weight is multiplied along): 1/4 for the 16 values of a patch, 1/8 for the 64 hidden values, 0.18 for the 32 channels. That’s what Chapter 10’s train.py did by hand.

A model’s name is a type. Functions take a model as a parameter and read its fields with ., so forward loses five parameters. Replace forward, loss and predict with:

## The class scores of each image.
fn forward(m: Spotter, img: Images u8) -> [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 |> head(m.w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, img), y, over=C)

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

m.tw1 is a [L N Dt], as declared, and m.tw1[l] is block l’s [N Dt], as before. Delete predict_saved, and the import python torch it needed: Section 3 saves and loads whole models, by name, with every shape checked, which beats w[3] as [L D Dc].

A field the model doesn’t have is an error, with a suggestion:

dim D = 4

model M:
    param w: [D D] = zeros

fn f(m: M, x: [b D]) -> [b D] = x @ m.W

$ shapescript check typo.ss
error[E0901]: `M` has no field `W`
 --> typo.ss:6:39
  |
3 | model M:
  |       - `M` is declared here
...
6 | fn f(m: M, x: [b D]) -> [b D] = x @ m.W
  |                                       ^
  |
  = help: a field with a similar name exists: `w`

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

Spotter.init(seed=0) makes a spotter, from ShapeScript or from Python. Each field draws its random numbers from a generator of its own, seeded from the model’s seed and the field’s name, so the same seed always gives the same spotter, and adding a field doesn’t change the others. From Python, a model is an ordinary torch.nn.Module with every field registered, so state_dict(), .parameters() and the rest of PyTorch work with it:

from spotter import Spotter, predict
from shapes import draw

m = Spotter.init(seed=0)
print(sum(p.numel() for p in m.parameters()))
img, y = draw(4)
print(predict(m, img))

$ shapescript run peek.py
12928
tensor([2, 2, 0, 2])

12,928 numbers, and untrained, it calls almost everything a plus. The rest of the chapter trains it.

Buffers

A model can also hold tensors that are saved with it but never trained: buffers, declared with buffer instead of param. Running statistics are the usual example. The spotter doesn’t need one, but a model that subtracts the average image from its inputs would keep that average as a buffer:

dim S = 16

model Centered:
    ## Subtracted from every image. Saved with the model, never trained.
    buffer center: [1 S S] = zeros

fn centered(m: Centered, img: [b 1 S S]) -> [b 1 S S] = img - m.center

Changing a model with mut

Functions so far have only read the spotter. To teach it a new set of shapes, you might keep the mixer blocks, which have learned to see strokes and corners, and start the head again from zero. That’s a change to the model, and a function can change a model only if its parameter is marked mut. Every call writes mut too, so the signature and the call both show what changes:

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

Forgetting either mut is an error:

model Spotter:
    param w_h: [D C] = normal(std=0.18)

fn reset_head(m: Spotter):
    m.w_h = m.w_h * 0.0

fn reset_head2(mut m: Spotter):
    m.w_h = m.w_h * 0.0

fn retrain(m: Spotter):
    reset_head2(m)

$ shapescript check no_mut.ss
error[E1102]: `m` can't be changed in `reset_head`
 --> no_mut.ss:5:5
  |
4 | fn reset_head(m: Spotter):
  |               - `m` is declared here
5 |     m.w_h = m.w_h * 0.0
  |     ^ `m` isn't a `mut` parameter
  |
  = help: to change it, declare it `mut m` and pass it as `mut` at every call

error[E1103]: `reset_head2` changes `m`, so its argument needs `mut`
  --> no_mut.ss:11:17
   |
 7 | fn reset_head2(mut m: Spotter):
   |                -------------- `m` is declared `mut` here
...
11 |     reset_head2(m)
   |                 ^
   |
   = help: write `mut m`

For more information about E1102 and E1103, run `shapescript explain` with the code.

Only models and optimizer state can be mut. Tensors and numbers are values: a function that computes a new one returns it.

Gradients

value_and_grad(f, m, args...) calls f(m, args...) and returns two things: its value, and its gradient with respect to m’s params:

## The loss and its gradient.
fn grads(m: Spotter, img: Images u8, y: [b] Index[C]) -> Grad[Spotter]:
    l, g = value_and_grad(loss, m, img, y)
    return g

g has type Grad[Spotter]: a record with a field for each param, g.tw1: [L N Dt] and so on (and none for buffers, which aren’t trained). Gradients of one model can be added, subtracted and scaled, and norm(g) is the norm of all their values together.

The checker reads f to see which params it uses. A loss that doesn’t use the model at all would always have a zero gradient, so it’s an error; a param it never reads gets a warning, which is often a forgotten layer. Say a forward mixes the channels in each block, but forgets to mix the patches:

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)

fn forward(m: Spotter, img: Images u8) -> [b C] Logit over C:
    x = embed(img |> to Unit |> to Signed, m.w_e)
    for l in 0..L:
        x = x |> mixer.mix_channels(m.cw1[l], m.cw2[l])
    return x |> mixer.norm |> pool |> matmul(m.w_h)

fn loss(m: Spotter, img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, img), y, over=C)

fn grads(m: Spotter, img: Images u8, y: [b] Index[C]) -> Grad[Spotter]:
    l, g = value_and_grad(loss, m, img, y)
    return g

$ shapescript check forgot.ss
warning[W1101]: `loss` never reads the param `tw1`, so its gradient is always zero
  --> forgot.ss:19:27
   |
 3 |     param tw1: [L N Dt] = normal(std=0.25)
   |           --- `tw1` is declared here
...
19 |     l, g = value_and_grad(loss, m, img, y)
   |                           ^^^^
   |
   = help: remove the param, or make it a `buffer` if it isn't trained

warning[W1101]: `loss` never reads the param `tw2`, so its gradient is always zero
  --> forgot.ss:19:27
   |
 4 |     param tw2: [L Dt N] = normal(std=0.125)
   |           --- `tw2` is declared here
...
19 |     l, g = value_and_grad(loss, m, img, y)
   |                           ^^^^
   |
   = help: remove the param, or make it a `buffer` if it isn't trained

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

The program would run, and train, and do worse than it should, with nothing in the loss curve to say why. Here the checker says which weights are dead.

Optimizers

sgd(m, lr=...) and adamw(m, lr=...) make optimizer state for a model. Its type comes from the model’s: adamw(m, lr=3e-3) with m: Spotter is an AdamW[Spotter], which covers every param of Spotter, so none can be left out, as one could be from the list Chapter 10 gave torch.optim.Adam. update(mut m, mut opt, g) takes one step:

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

## `steps` steps on the same batch.
fn train(mut m: Spotter, img: Images u8, y: [b] Index[C], steps: int):
    opt = adamw(m, lr=warmup(20, cosine(3e-3, steps)))
    for i in 0..steps:
        l = step(mut m, mut opt, img, y)

• update changes both the model and the optimizer state, so both are mut.
• clip_norm(1.0) is a gradient transform: a function from Grad[Spotter] to Grad[Spotter], used with a pipe. It scales the gradient down when its norm is above 1, which keeps one bad batch from throwing the weights far. Your own transforms are ordinary functions, such as fn halve(g: Grad[Spotter]) -> Grad[Spotter] = g * 0.5.
• The learning rate can be a number or a schedule: warmup(20, then) rises from 0 over 20 steps, then follows then; cosine(peak, steps) decays from peak to 0. The optimizer state counts its own updates, so there’s no scheduler to step.

Optimizer state is matched to its model. Passing an AdamW[Other] to update a Spotter is an error, as is an SGD[Spotter] where an AdamW[Spotter] is expected.

train only ever sees one batch. Real training needs a dataset, shuffled into batches, which is the next section.

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)

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

## The class scores of each image.
fn head(x: [b n D], w: [D C]) -> [b C] Logit over C = pool(x) @ w

## 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) -> [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 |> head(m.w_h)

## The mean cross-entropy of the scores against the labels.
fn loss(m: Spotter, img: Images u8, y: [b] Index[C]) -> [] NonNeg:
    return cross_entropy(forward(m, img), y, over=C)

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

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

## The loss and its gradient.
fn grads(m: Spotter, img: Images u8, y: [b] Index[C]) -> Grad[Spotter]:
    l, g = value_and_grad(loss, m, img, y)
    return g

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

## `steps` steps on the same batch.
fn train(mut m: Spotter, img: Images u8, y: [b] Index[C], steps: int):
    opt = adamw(m, lr=warmup(20, cosine(3e-3, steps)))
    for i in 0..steps:
        l = step(mut m, mut opt, img, y)

## 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. Write sgd_head(mut m: Spotter, img: Images u8, y: [b] Index[C], lr: float), which updates only the head by hand, without an optimizer: w_h minus lr times its gradient, leaving the mixer blocks as they are. That’s how you’d retrain the head after reset_head.

2. Find the bugs:

   dim D = 32
   dim C = 4
   
   model Head:
       param w: [D C] = normal(std=0.18)
   
   fn loss(m: Head, x: [b D], y: [b] Index[C]) -> [] = cross_entropy(x @ m.w, y, over=C)
   
   fn step(m: Head, opt: SGD[Head], x: [b D], y: [b] Index[C]):
       l, g = value_and_grad(loss, m, x, y)
       update(m, opt, g)
   

3. Write a gradient transform scale_down(g: Grad[Spotter], k: float) -> Grad[Spotter], and use it in step’s pipe after clip_norm, to take steps half as long.