Checkpoints

save

save(path, values…) writes a checkpoint: the values, in order, to the file path, replacing it. The path is a string or a Python value, given first, by position (E1603); the values follow, by position (E0406 for a keyword or for no values). It has no result.

fn run(mut m: Mlp, mut opt: AdamW[Mlp], train: Dataset[Digit], first: int, epochs: int, mode: Mode):
    for epoch in first..epochs:
        for batch in train |> shuffle(seed=epoch) |> batches(64):
            l, g = value_and_grad(loss, m, batch, mode)
            update(mut m, mut opt, g)
        save("ckpt.pt", m, opt, epoch + 1)

fn resume(train: Dataset[Digit], epochs: int, mode: Mode) -> Mlp:
    m, opt, done = load("ckpt.pt", Mlp, AdamW[Mlp], int)
    run(mut m, mut opt, train, done, epochs, mode)
    return m

A checkpoint holds models (every field, params and buffers), optimizer state (everything state_dict() gives: the optimizer’s moments or momentum, its hyperparameters, its learning rate or schedule, and its update count) and host scalars, int, float and bool, such as an epoch or step count. Anything else is E1601. A tensor can’t be saved on its own: tensor state that training needs belongs in a model, as a buffer if it isn’t trained, so that its type, and the file’s, say what it is. A Python value is accepted and checked when save runs, which raises TypeError for anything a checkpoint can’t hold.

load

load(path, Types…) reads a checkpoint. After the path come the types of the values in the file, in order, written as in a signature: a model’s name (Mlp), optimizer state (AdamW[Mlp], SGD[Mlp]), or int, float or bool. Its result is a value of each type: one type gives one value, and more give results to name, as m, opt, done = load("ckpt.pt", Mlp, AdamW[Mlp], int); a different number of names is E0411. A value, a tensor type, or any other type, such as Grad[Mlp], a Schedule, a Mode or a logs record, is E1602, as is an expression that isn’t a type; a name that isn’t declared is E0301, and AdamW without a model, or of something other than a model, is E0303. The types are given by position (E0406), and at least one is (E0406).

What load makes is new: a model and optimizer state it loads belong to the function, which can pass them as mut (Changing models with mut).

Checkpoints at run time

save and load compile to rt.save(path, values…, where=…) and rt.load(path, types…, where=…), with where the call’s file:line:col; from Python they’re shapescript.save and shapescript.load, whose types are model classes, shapescript.AdamW[Mlp] and shapescript.SGD[Mlp], and int, float and bool.

The file is written with torch.save, to a file beside path that is then renamed to it, so a save that is interrupted leaves the last checkpoint whole. It’s a dict, {"format": "shapescript.checkpoint", "version": 1, "values": [...]}, with one entry per value:

• a model: {"kind": "model", "type": "Mlp", "fields": [...], "tensors": {"w1": tensor, ...}};
• optimizer state: {"kind": "optimizer", "type": "AdamW[Mlp]", "optimizer": "adamw", "model": "Mlp", "fields": [...], "of": 0, "state": {...}}, where state is the state’s state_dict(), fields are the params it was updating, and of is the position of the model saved with it whose params those were, or None;
• a host scalar: {"kind": "int", "value": 3}, and the same with "float" and "bool".

Each field is {"name": "w1", "kind": "param", "shape": [784, 256], "dtype": "float32"}.

Loading reads the file with torch.load(path, map_location=device, weights_only=True), so the file can’t run code, and every tensor goes straight onto the program’s declared device, never the device it was saved from: a checkpoint saved on a GPU loads on a CPU, and the other way round. Before it makes anything, load compares the file with the types, and raises shapescript.ShapeError (a ValueError) listing everything that differs: each field whose shape or dtype differs, with both, each field one has and the other doesn’t, and each value of another kind or type. Loading an Mlp saved with Hd = 128 as one with Hd = 256 gives:

ShapeError: train.ss:9:20: `ckpt.pt` doesn't match the types it's loaded as:
  value 1, `Mlp`: field `w1` is (784, 128) float32 in the file, but `Mlp` declares (784, 256) float32
  value 1, `Mlp`: field `w2` is (128, 10) float32 in the file, but `Mlp` declares (256, 10) float32
  value 2, `AdamW[Mlp]`: field `w1` is (784, 128) float32 in the file, but `Mlp` declares (784, 256) float32
  value 2, `AdamW[Mlp]`: field `w2` is (128, 10) float32 in the file, but `Mlp` declares (256, 10) float32

A file with another number of values is a ShapeError naming both lists (`ckpt.pt` holds 3 values (Mlp, AdamW[Mlp], int), but `load` was given 2 types (Mlp, AdamW[Mlp])), and so is a file save didn’t write. A model is compared by its name and fields, as a model parameter is, and an int in the file can be loaded as a float.

A model is made with the file’s tensors, without drawing new ones. Optimizer state is made with the file’s hyperparameters, learning rate and update count, so a schedule carries on where it was; saved with its model, it’s bound to the loaded model, and otherwise to stand-in params until update moves it to the model it’s given. Training that saves after n steps, loads, and carries on gives the same result as training that didn’t stop (tests/python/test_checkpoints.py checks it bit for bit, dropout and a schedule included). Random draws need the same keys in both runs for that, and they have them: a mode’s draws inside a loop are keyed by the iteration, so a resumed loop over done..epochs draws what the uninterrupted one draws in those epochs.