Training loops

There are no hooks, callbacks or trainer objects: a training loop is a function the program writes, from value_and_grad, update, datasets, logs and checkpoints.

The README’s training example

examples/train_mlp.ss is the README’s training example, exactly (a test compares them): a two-layer classifier, train_step with clipping and logs, and fit(mut m, train, sink, epochs, mode), which makes AdamW state and runs epochs epochs, shuffling each with its number, writing a record per step, flushing the sink and saving m, the state and the number of epochs done after each. resume(train, sink, epochs, mode) loads that checkpoint and carries on up to epochs. Called from Python with mode=shapescript.Eval, so that dropout is off, its losses for 100 steps match a hand-written PyTorch loop’s (torch.optim.AdamW, clip_grad_norm_) to within 1e-5, on the CPU and on a GPU, and stopping after 5 epochs and resuming gives the 100 steps of a run that didn’t stop (tests/python/test_train_mlp.py). The mode is passed to every train_step, so with Train(seed=s) each step draws its own masks, and a resumed run draws new ones, as above.

The standard library’s fit

The standard library’s fit is a generic function, in shapescript/std/fit.ss in the Python package (and built into the compiler): fit(mut m: M, train: Dataset[E], sink, epochs, size, lr, mode) trains any model M on any data record E for which the program declares loss(m, batch, mode) -> [], next to the model; it makes AdamW state and runs run, which shuffles each epoch with its number, calls train_step (AdamW, the gradient’s norm clipped to 1, a StepLogs record of the loss and the gradient’s norm per step), flushes the sink and saves m, the state and the epochs done to ckpt.pt after each epoch. resume(like: M, train, sink, epochs, size, mode) carries on from that checkpoint; like only says which model to load. Import it and call it from ShapeScript:

import shapescript.std.fit as std

fn train(mut m: Mlp, data: Dataset[Digit], sink: LogSink, epochs: int, mode: Mode):
    std.fit(mut m, data, sink, epochs, 64, 3e-4, mode)

A model without the loss it needs is E2311 at the call. It has no hooks or callbacks, and it’s short enough to read on one screen; to change the loop, copy the file (shapescript.std.path("fit") gives its path). With size=64 and lr=3e-4 it gives the README’s fit’s losses on the README’s model, and the same as a fit written for a second model, with its own forward and loss (tests/python/test_generics.py).