3.9. Saving and Loading Weights

RustyML persists a trained Sequential model with 3 methods: save_to_path, load_from_path, and the opt-in lenient load_partial_from_path. These methods are narrow by design. RustyML saves the layer weights alone, and the file does not carry enough information to rebuild a model. Loading takes 2 steps. First, collect the identical layer stack in a SequentialBuilder and build it against the same input shape. Second, load the saved arrays into that built model.

This section describes what the file stores, why the boundary sits there, the failure modes, and workflow recipes for a weights-only model.

3.9.1. What actually gets written to disk

save_to_path walks the layers. For each layer, it records the layer type name and every named array the layer holds. It then serializes the whole checkpoint with postcard into a compact binary blob, and a buffered writer writes that blob to disk. Nothing is copied on the way, because each record borrows the live array until postcard reads it. load_from_path reads the file back, deserializes it, checks that the built model matches the file, and applies the arrays layer by layer. The signatures:

pub fn save_to_path(&self, path: impl AsRef<std::path::Path>) -> RustymlResult<()>;
pub fn load_from_path(&mut self, path: impl AsRef<std::path::Path>) -> RustymlResult<()>;
pub fn load_partial_from_path(&mut self, path: impl AsRef<std::path::Path>) -> RustymlResult<LoadReport>;

All 3 methods accept any type that implements AsRef<Path>: &str, String, Path, or PathBuf. The "model.bin" literals in the examples below are only the common case. The file extension has no effect on the format, because postcard writes raw bytes whatever the file name is. save_to_path uses File::create, which truncates and overwrites the file, so a second save to the same path replaces the previous checkpoint. This behavior fits a “keep only the best” loop.

The per-layer metadata serves 1 purpose: validation, and not reconstruction. Each layer contributes a type-name string, for example "Dense" or "Conv2D". Each array contributes its own name and its kind. These are validation tags, and not a recipe for building a layer. The file does not record the activation function of a layer.

It does not record epsilon, momentum, stride, kernel size, dilation, or group count. An empty Sequential handed to load_from_path does not come back as a working model. RustyML calls this format “weights-only” for that reason. The architecture lives in the source code, and the file carries only the numbers that fill it.

The file also holds the shapes each layer was built for. A layer allocates its arrays in build, from the shapes of its inputs, so those shapes decide every extent the layer holds. The build record holds 1 shape per input of the layer. A layer with 1 input therefore records 1 shape, and a merge layer with 3 inputs records 3. Recording them lets a load refuse a model that was built for another input, and the refusal names the layer and both sides.

The batch axis of that record is free. A layer serves every batch size, so a checkpoint written by a model built for 1 sample loads into a model built for 32. A layer that owns no array and reads no extent of its input carries no build record at all. An activation and Dropout are such layers, and a load skips the comparison for them.

3.9.2. The named checkpoint: 1 address per array

The file addresses every array by a dotted path, <scope>.<name>. scope is the position of the layer, counted from the input. name is the name the layer gives the array. The kernel of the first layer of a model is therefore 0.kernel. The moving mean of a BatchNormalization at position 1 is 1.moving_mean. This is the same pair the optimizer keys its state on, so 1 address serves the training loop and the file alike (see 3.4. Optimizers).

A Graph model writes the same format under the same paths. Its scope is the position of the layer in the layer arena, and not the position of a node. A layer that several nodes call therefore holds 1 set of paths, whatever number of nodes read it.

Each record also carries a WeightKind. Trainable says an optimizer updates the array, and it covers every kernel, every bias, and every normalization scale and shift. NonTrainable says the layer keeps the array and no optimizer writes it. The 2 moving statistics of BatchNormalization are the only non-trainable arrays in the crate today. A load compares the kind next to the shape. A file that offers a trainable array where the layer keeps state is therefore refused, and not applied.

weight_paths() lists every address of a model, in the order a file holds them. weight(path) reads 1 of them, as a view that borrows the live array:

use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::activation::linear::Linear;
use rustyml::neural_network::layers::dense::Dense;
use rustyml::neural_network::layers::regularization::normalization::batch_normalization::BatchNormalization;
use rustyml::neural_network::sequential::SequentialBuilder;

fn main() {
    let model = SequentialBuilder::new()
        .add(Dense::new(3, Linear::new()).unwrap())
        .add(BatchNormalization::new(0.9, 1e-5).unwrap())
        .add(Dense::new(2, Linear::new()).unwrap().with_use_bias(false))
        .build(&Shape::with_free_batch(&[2, 4]))
        .unwrap();

    // Every array of the model, under the address a saved file gives it.
    for path in model.weight_paths() {
        println!("{path:<18} {:?}", model.weight(&path).unwrap().shape());
    }
}

0.kernel           [4, 3]
0.bias             [1, 3]
1.gamma            [3]
1.beta             [3]
1.moving_mean      [3]
1.moving_variance  [3]
2.kernel           [3, 2]

The third layer of that model shows what an optional array does to the list. with_use_bias(false) takes 2.bias out of the model, so it is out of the file as well. The remaining paths do not move, because a path names an array and never a position inside a layer. with_center(false) and with_scale(false) do the same to beta and gamma on the 4 normalization layers.

The arrays each layer holds:

This difference among the normalization layers is correct by design. BatchNormalization accumulates moving statistics during training and uses them at inference. Those 2 arrays are part of the trained state, so they must survive serialization. If a reload dropped them, the model would normalize eval-mode inputs with default statistics and produce wrong output.

Layer, instance, and group normalization compute their statistics from the current input on every forward pass, so they hold no moving state. gamma and beta are all these layers need to save. 3.9.4 shows the BatchNormalization case in detail.

3.9.3. A complete round-trip

The full loop has 5 steps: build, train briefly, save, rebuild the same stack, and load. The last step confirms that the restored model predicts the same values. Note the make_arch function. It defines the architecture once, the input shape included, and both the live model and the reload target call it. This habit is the most useful practice for weights-only persistence, because it guarantees that the 2 stacks cannot drift apart. The build shapes travel in the file as well, so the load, and not the first prediction, catches a drift there.

use ndarray::Array;
use rustyml::error::Error;
use rustyml::neural_network::Tensor;
use rustyml::neural_network::layers::activation::linear::Linear;
use rustyml::neural_network::layers::dense::Dense;
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::Shape;
use rustyml::neural_network::sequential::{Sequential, SequentialBuilder};

// Define the architecture once, input shape included. Reuse it for the live model and the
// reload target.
fn make_arch() -> Sequential {
    SequentialBuilder::new()
        .add(Dense::new(3, Linear::new()).unwrap())
        .add(Dense::new(2, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(&[2, 4]))
        .unwrap()
}

fn main() -> Result<(), Error> {
    let x: Tensor = Array::from_shape_vec((2, 4), vec![0.1f32, 0.2, 0.3, 0.4, -0.1, -0.2, -0.3, -0.4])
        .unwrap()
        .into_dyn();
    let y: Tensor = Array::from_shape_vec((2, 2), vec![1.0f32, 0.0, 0.0, 1.0])
        .unwrap()
        .into_dyn();

    let mut model = make_arch();
    model.compile(SGD::new(0.01, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());
    model.fit(&x, &y, 5)?;
    let before = model.predict(&x)?;

    // Save weights, then rebuild the identical stack and load into it.
    let path = "roundtrip_demo.bin";
    model.save_to_path(path)?;

    let mut restored = make_arch();
    restored.load_from_path(path)?;
    let after = restored.predict(&x)?;

    // The two prediction tensors must agree element-wise.
    let max_diff = (&after - &before)
        .mapv(f32::abs)
        .iter()
        .cloned()
        .fold(0.0f32, f32::max);
    println!("max abs difference after round-trip: {max_diff:e}");
    assert!(max_diff < 1e-6);

    std::fs::remove_file(path).unwrap();
    Ok(())
}

The round-trip is exact, and not approximate. postcard stores each f32 losslessly. load_from_path writes the arrays straight into the layers, so the reloaded model runs the same computation, bit-for-bit. The 1e-6 tolerance in the example is defensive slack, and not a hedge against drift.

3.9.4. Normalization running statistics survive the round-trip

The inference path of BatchNormalization reads moving_mean and moving_variance. This example trains those statistics away from their defaults, saves the model, and reloads the weights into a fresh, untrained model. It then checks that eval-mode predictions still match. They do. The 2 arrays travel under their own paths, marked NonTrainable, next to the trainable gamma and beta of the same layer.

use ndarray::Array;
use rustyml::neural_network::Tensor;
use rustyml::neural_network::layers::regularization::normalization::batch_normalization::BatchNormalization;
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::Shape;
use rustyml::neural_network::sequential::{Sequential, SequentialBuilder};

fn make_arch() -> Sequential {
    SequentialBuilder::new()
        .add(BatchNormalization::new(0.9, 1e-5).unwrap())
        .build(&Shape::with_free_batch(&[4, 3]))
        .unwrap()
}

fn main() {
    let x: Tensor = Array::from_shape_vec(
        (4, 3),
        vec![0.5f32, -1.0, 2.0, 1.5, 0.2, -0.7, -1.2, 0.8, 1.1, 0.3, -0.4, 0.9],
    )
    .unwrap()
    .into_dyn();

    // Train so moving_mean / moving_variance move away from their initial values.
    let mut model = make_arch();
    model.compile(SGD::new(0.001, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());
    model.fit(&x, &x, 8).unwrap();
    let before = model.predict(&x).unwrap(); // eval mode uses the running stats

    let path = "batchnorm_demo.bin";
    model.save_to_path(path).unwrap();

    let mut restored = make_arch(); // fresh: running stats at their defaults
    restored.load_from_path(path).unwrap();
    let after = restored.predict(&x).unwrap();

    let max_diff = (&after - &before)
        .mapv(f32::abs)
        .iter()
        .cloned()
        .fold(0.0f32, f32::max);
    println!("running-stat round-trip max abs difference: {max_diff:e}");
    assert!(max_diff < 1e-6);

    std::fs::remove_file(path).unwrap();
}

If the file did not save the moving statistics, after would differ sharply from before. The defaults of the fresh model do not match the 8 epochs of accumulated batch statistics.

3.9.5. What is not saved, and why it bites

The optimizer, its accumulated state, the loss function, and the entire compile configuration stay behind. Loading gives weights on a blank model. The optimizer and loss fields are None, so 2 consequences follow.

First, the model needs another compile call before it can fit or train. Prediction works right away, because predict needs no optimizer. Second, a resumed training run after a load restarts the optimizer from zero. This point is easy to miss. The first-moment and second-moment estimates of Adam reset to zero.

The momentum buffer of SGD resets too. The bias-correction timestep of Adam starts over at 1. For a few fine-tuning steps, this reset causes no harm. For a long training run split across a save and load, the first post-load steps take larger, less-damped updates than an uninterrupted run would take. The loss curve shows a temporary bump.

This bump is a persistence artifact, and not a data problem. To pause and resume long training without this discontinuity, keep the process alive. Do not pause it with a save to disk and a reload. RustyML has no API to serialize optimizer moments. The fit-time shuffle seed is also not saved. Set it again with set_seed after a load, for a reproducible resumed shuffle.

3.9.6. Error variants on load

Every failure surfaces as Error::Io(...) (see 1.6. Error Handling). There are exactly 4 variants, and each one maps to a distinct cause of failure:

The load path checks the header first, so these 4 errors follow an order. A file that is not a model at all never reaches the postcard decoder. A file from an incompatible release never reaches the structural checks.

ModelStructureMismatch is the error a model hits most often while it is under development. The strict load runs 6 checks, and it runs them in this order. First the layer count, then the layer type of each position. Then the build shapes, where both sides carry them. Then the number of arrays of a layer, and the name of each one.

Then the kind of each array. Last the shape and the element count of each array. The whole first pass reads the model alone and writes nothing. A refusal therefore leaves the model exactly as it was. A model never holds the arrays of 1 file next to the arrays of another.

The message names the checkpoint path, or the layer position for a whole-layer disagreement. The build-shape check usually fires first, and it fires before it compares any array. Take a Dense model built for (None, 4), loaded from a file that a model built for (None, 5) wrote. It fails with model structure mismatch: layer 0 (`Dense`) was built for input shape (None, 4), and the file records (None, 5). A merge layer names 1 shape per input on both sides of that message. Matching the error is straightforward:

use rustyml::error::{Error, IoError};
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::activation::linear::Linear;
use rustyml::neural_network::layers::dense::Dense;
use rustyml::neural_network::sequential::SequentialBuilder;

fn main() {
    // Save a 1-layer model.
    let saved = SequentialBuilder::new()
        .add(Dense::new(2, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(&[2, 2]))
        .unwrap();
    let path = "mismatch_demo.bin";
    saved.save_to_path(path).unwrap();

    // Rebuild with the wrong layer count and try to load it.
    let mut wrong = SequentialBuilder::new()
        .add(Dense::new(2, Linear::new()).unwrap())
        .add(Dense::new(2, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(&[2, 2]))
        .unwrap();

    match wrong.load_from_path(path) {
        Err(Error::Io(IoError::ModelStructureMismatch(msg))) => {
            println!("rejected as expected: {msg}");
        }
        Err(other) => panic!("unexpected error: {other:?}"),
        Ok(()) => panic!("load must not succeed on a structure mismatch"),
    }

    std::fs::remove_file(path).unwrap();
}

The type-name check compares strings. It catches common mistakes: a Conv layer where the file has a Dense layer, or a missing or extra layer. It does not catch a hyperparameter change that leaves the type name and weight shapes unchanged. 2 Dense layers with the same shapes but different activations load without complaint, because the activation is not on disk. This gap is the direct cost of storing metadata tags instead of a full architecture, and the reason the make_arch discipline matters.

Strict is the default, and the type name is why. A name and a shape together are a weaker key than a closed set of per-layer containers. InstanceNormalization, GroupNormalization, and a rank-2 LayerNormalization all expose gamma [C] and beta [C], identically. Nothing but the per-layer type name tells the 3 apart. A lenient default would therefore load a file that is wrong for the model and apply every array of it.

The difference would surface later, as a wrong prediction. So the load refuses on any disagreement, and the lenient path is the separate load_partial_from_path, which a caller asks for by name. 3.9.9 covers it.

3.9.7. The postcard format: size, speed, portability

postcard is a compact, non-self-describing binary format. Non-self-describing means the stream carries no field names. It writes only the values, in declaration order. This design keeps the files small. It also ties each file to the exact struct and enum layout of the crate version that wrote it. File size is easy to predict.

Each parameter is an f32, 4 bytes. A small fixed overhead adds the array shape (varint-encoded) and 1 byte for the kind of each array. It also adds the build shapes and the short strings: 1 layer type name per layer, and 1 array name per array. The Total params: N (N*4 B) line that Sequential::summary prints is therefore a close upper estimate of the file size. A 100,000-parameter model takes about 400 KB. Serialization runs as a single linear pass with 1 buffered write, so I/O time, and not CPU time, dominates the cost.

The format is portable across machines. postcard defines its own byte order instead of a dump of native-endian memory. A checkpoint written on 1 architecture therefore deserializes correctly on another, whatever the endianness of the host. This portability needs no conversion step and no per-platform variant. The format is not portable across crate versions. That limit is the subject of the next section.

3.9.8. Workflow recipes

Checkpoint the best model. Train in short rounds. Score the model with evaluate after each round, and overwrite a single file whenever the score improves. save_to_path truncates the file, so it always holds the best weights found so far. The final in-memory model may have overfit past the best point, so discard it and use the reload instead. Score with evaluate, and not with the last entry of the History that fit returns.

A history entry is the loss measured during the epoch, on a forward pass taken before the weight update of that epoch. The last entry therefore describes weights the model no longer holds, and a rule based on it would checkpoint a round late. evaluate runs 1 inference-mode forward pass over the data and scores it with the compiled loss. It updates nothing: no gradients, no parameters, and no BatchNormalization running statistics. A selection rule built on evaluate therefore cannot change the training it measures.

use ndarray::Array;
use rustyml::error::Error;
use rustyml::neural_network::Tensor;
use rustyml::neural_network::layers::activation::linear::Linear;
use rustyml::neural_network::layers::dense::Dense;
use rustyml::neural_network::losses::MeanSquaredError;
use rustyml::neural_network::optimizers::SGD;
use rustyml::neural_network::Shape;
use rustyml::neural_network::sequential::{Sequential, SequentialBuilder};

fn make_arch() -> Sequential {
    SequentialBuilder::new()
        .add(Dense::new(3, Linear::new()).unwrap())
        .add(Dense::new(2, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(&[2, 4]))
        .unwrap()
}

fn main() -> Result<(), Error> {
    let x: Tensor = Array::from_shape_vec((2, 4), vec![0.1f32, 0.2, 0.3, 0.4, -0.1, -0.2, -0.3, -0.4])
        .unwrap()
        .into_dyn();
    let y: Tensor = Array::from_shape_vec((2, 2), vec![1.0f32, 0.0, 0.0, 1.0])
        .unwrap()
        .into_dyn();

    let mut model = make_arch();
    model.compile(SGD::new(0.05, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());

    let path = "best.bin";
    let mut best = f32::INFINITY;
    for round in 0..10 {
        model.fit(&x, &y, 2)?;
        let val = model.evaluate(&x, &y)?;
        if val < best {
            best = val;
            model.save_to_path(path)?; // overwrites the previous checkpoint
            println!("round {round}: new best {val:.6}, checkpoint written");
        }
    }

    // best.bin holds the lowest-loss weights, not necessarily the final ones.
    let mut deployed = make_arch();
    deployed.load_from_path(path)?;
    // `evaluate` needs the compiled loss. It never touches the optimizer passed here.
    deployed.compile(SGD::new(0.05, 0.0, false, 0.0).unwrap(), MeanSquaredError::new());
    assert_eq!(deployed.evaluate(&x, &y)?, best);

    std::fs::remove_file(path).unwrap();
    Ok(())
}

The final check uses exact equality, and not a tolerance. The reloaded weights are bit-for-bit the same as the saved ones, and evaluate is deterministic on a model with no dropout. The score of the restored checkpoint must therefore equal the score that caused RustyML to write it.

Transfer weights between programs. A training binary builds the architecture, trains the model, and calls save_to_path. A separate serving binary builds the identical architecture and calls load_from_path. Share the make_arch function through a common module or crate. The 2 programs then cannot disagree on the architecture, or on the input shape it builds for.

The type-name, count, and build-shape checks catch a drift, but only after a failed load. A shared constructor catches the same drift at compile time. The serving binary never needs to call compile, because it only calls predict.

Manage versions across crate upgrades. Every saved model opens with a magic tag and a format version, and load_from_path validates both before it decodes anything else. A checkpoint written by a release with a different weight layout therefore fails right away with IoError::UnsupportedModelFormat. Without this check, the load could parse the file into arrays that look correct but hold wrong values. The structural checks that run after the header compare layer counts, type names, and build shapes.

They compare the name, the kind, and the shape of every array. A stale file can satisfy every one of them by coincidence. A square convolution kernel, for example, keeps its extents when its axes are permuted. A Dense weight shape stays the same whatever tensor layout produced the input feeding it.

This guard depends on developer discipline. MODEL_FORMAT_VERSION is 3, and version 3 is the named checkpoint of 3.9.2. Its build record holds 1 shape per input of a layer. A merge layer with 3 inputs therefore records 3 shapes, and a layer with 1 input records 1. Every earlier version stops loading, and the refusal names both version numbers. Re-save such a checkpoint from a model built under this release.

RustyML bumps the version on any change to the layout of a record. It also bumps on a change to the order of the fields of a structure, or to the meaning of a field. A new layer needs no bump. A new layer adds a type name and some array names, and both travel in the file as strings.

The load catches a change that gets a version bump. It does not catch a change that nobody bumped, because postcard is still non-self-describing and nothing else in the file is labeled. For a checkpoint that must reload weeks or versions later, pin the RustyML version in Cargo.toml. After a deliberate upgrade, run training again, or load under the old version and save again under the new one. Do not trust an old file against new code. The header makes a version failure loud and immediate instead of silent.

3.9.9. Loading part of a file on purpose

load_partial_from_path is the lenient load, and a caller asks for it by name. It writes an array under 3 conditions. The position holds the same layer type, the 2 sides agree on the build shapes, and the file holds the same name, kind, and shape. Nothing about the layer roster fails. It returns a LoadReport with 3 lists of checkpoint paths.

applied holds what took a value, and missing holds the paths of the model that got none. unused holds the paths of the file that reached no array. An array whose shape or kind disagrees is in missing and in unused at the same time. The model got no value for it, and the file value went nowhere. A position whose layer type or build shape differs puts every path of that layer in both lists. A name and a shape cannot tell 2 normalization layers apart.

The header is still checked. A file of another format version carries bytes that mean something else, so it is an error here as well.

use rustyml::error::Error;
use rustyml::neural_network::Shape;
use rustyml::neural_network::layers::activation::linear::Linear;
use rustyml::neural_network::layers::dense::Dense;
use rustyml::neural_network::sequential::SequentialBuilder;

fn main() -> Result<(), Error> {
    let saved = SequentialBuilder::new()
        .add(Dense::new(3, Linear::new()).unwrap())
        .add(Dense::new(2, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(&[2, 4]))
        .unwrap();
    let path = "partial_demo.bin";
    saved.save_to_path(path)?;

    // The target keeps the first layer and widens the second one.
    let mut target = SequentialBuilder::new()
        .add(Dense::new(3, Linear::new()).unwrap())
        .add(Dense::new(5, Linear::new()).unwrap())
        .build(&Shape::with_free_batch(&[2, 4]))
        .unwrap();

    // The strict load refuses, and the message names the path that disagrees.
    let refusal = target.load_from_path(path).unwrap_err();
    println!("strict: {refusal}");

    // The lenient load takes what matches and reports the rest.
    let report = target.load_partial_from_path(path)?;
    println!("applied: {:?}", report.applied);
    println!("missing: {:?}", report.missing);
    println!("unused:  {:?}", report.unused);

    std::fs::remove_file(path).unwrap();
    Ok(())
}

strict: model structure mismatch: `1.kernel` has shape [3, 5] in the model, and [3, 2] in the file
applied: ["0.kernel", "0.bias"]
missing: ["1.kernel", "1.bias"]
unused:  ["1.kernel", "1.bias"]

This is the transfer-learning path. Load a trained backbone into a model whose head is new or resized. Then read missing to confirm that only the head went unfilled. Read unused to confirm that the file gave up nothing worth keeping.

Treat a non-empty missing list as a result to check, and never as a warning to skip. The strict load stays the right choice everywhere else, a restart of a training run included. There the file and the model must agree completely.

For more detail on the persistence machinery, see 7.2. Model Persistence in Depth. That section covers the record layout of the file, the 2-pass load, and how the same postcard format serves the classical estimators.