Keyboard shortcuts

Press or to navigate between chapters

Press S or / to search in the book

Press ? to show this help

Press Esc to hide this help

3.9. Saving and Loading Weights

RustyML persists a trained Sequential model with exactly 2 methods: save_to_path and load_from_path. These methods are narrow by design. Keras’ model.save() writes a self-describing bundle that rebuilds the graph, the compiler config, and the optimizer state. RustyML does not do this. RustyML saves only the layer weights, and the file does not carry enough information to rebuild a model. Loading takes 2 steps. First, you build the identical layer stack in code. Second, you load the saved arrays into that stack. 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 a small metadata tag plus the layer’s weights. It then serializes the whole vector with postcard into a compact binary blob, and a buffered writer writes that blob to disk. load_from_path reads the file back, deserializes it, checks that the model you built 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<()>;

Both 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, since postcard writes raw bytes no matter what name you give the file. save_to_path uses File::create, which truncates and overwrites the file, so saving again to the same path replaces the previous checkpoint. This behavior fits a “keep only the best” loop.

The per-layer metadata serves one purpose: validation, not reconstruction. Each layer contributes a type-name string, for example "Dense" or "Conv2D", and an output-shape string. These strings are validation tags, not a recipe for building a layer. The file does not record a layer’s activation function. It does not record epsilon, momentum, stride, kernel size, dilation, or group count. You cannot hand load_from_path an empty Sequential and get back a working model. RustyML calls this format “weights-only” for that reason. The architecture lives in your source code, and the file carries only the numbers that fill it.

Persisted in the fileNot persisted
Per-layer type-name tag (validated on load)Activation functions, hyperparameters (epsilon, momentum, stride, …)
Per-layer output-shape tag (informational only)Optimizer and its accumulated state (Adam moments, SGD momentum)
Every layer’s weight arrays (see 3.9.2)Loss function and compile state
BatchNormalization running mean and varianceFit-time shuffle seed, training-mode flag

3.9.2. The LayerWeight enum: per-layer payloads

The on-disk weight format is a Rust enum, LayerWeight<'a>. It has one variant for each supported layer type, plus an Empty variant for layers with no parameters. Each variant wraps a small struct that stores its arrays as Cow. Cow lets one type serve both directions. When saving, Sequential::get_weights borrows the live arrays (Cow::Borrowed), with no clone. When loading, load_from_path deserializes into owned arrays (Cow::Owned). The enum uses serde’s default, externally tagged representation, because postcard is not self-describing and needs the discriminant written out explicitly.

Each variant’s payload determines whether a round-trip is exact:

VariantPayload
Denseweight (in, out), bias (1, out)
SimpleRNNkernel, recurrent_kernel, bias
LSTM / GRUfused kernel, recurrent_kernel, bias (gate blocks [i|f|g|o] / [z|r|h])
Conv1D / Conv2D / Conv3Dconvolution weight kernel and bias
SeparableConv2Ddepthwise weight kernel, pointwise weight kernel, and bias
DepthwiseConv2Ddepthwise weight kernel and bias (no pointwise kernel)
BatchNormalizationgamma, beta, running_mean, running_var
LayerNormalization / InstanceNormalization / GroupNormalizationgamma, beta only
Emptynothing (Dropout, pooling, flatten, pure activation layers)

This difference among the normalization layers is correct by design. BatchNormalization accumulates running 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 running 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, 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 two stacks cannot drift apart.

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::sequential::Sequential;

// Define the architecture once. Reuse it for the live model and the reload target.
fn make_arch() -> Sequential {
    let mut m = Sequential::new();
    m.add(Dense::new(4, 3, Linear::new()).unwrap())
        .add(Dense::new(3, 2, Linear::new()).unwrap());
    m
}

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, not approximate. postcard stores each f32 losslessly, and 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, not a hedge against drift.

3.9.4. Normalization running statistics survive the round-trip

BatchNormalization’s inference path reads running_mean and running_var. This test 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 running arrays travel with the BatchNormalization variant.

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::sequential::Sequential;

fn make_arch() -> Sequential {
    let mut m = Sequential::new();
    m.add(BatchNormalization::new(vec![4, 3], 0.9, 1e-5).unwrap());
    m
}

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 running_mean / running_var 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 running statistics, after would differ sharply from before. The fresh model’s defaults 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. This is the biggest difference from Keras. Keras’ default save format bundles the optimizer, so fit can resume training without a gap. In RustyML, loading gives you weights on a blank model. The optimizer and loss fields are None, so 2 consequences follow.

First, you must call compile again before the model can fit or train. Prediction works right away, because predict needs no optimizer. Second, resuming training after a load restarts the optimizer from zero. This point is easy to miss. Adam’s first-moment and second-moment estimates reset to zero. SGD’s momentum buffer resets too. Adam’s bias-correction timestep 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, not a data problem. To pause and resume long training without this discontinuity, keep the process alive instead of saving to disk and reloading. RustyML has no API to serialize optimizer moments. The fit-time shuffle seed is also not saved. Re-set it with set_seed after loading, if you want the resumed shuffle to be reproducible.

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:

ErrorCause
IoError::StdThe file could not be opened/read (missing path, permissions) or written
IoError::UnsupportedModelFormatThe file does not carry this build’s magic tag and format version. It is not a RustyML model, or a release with a different on-disk weight layout wrote it (see 3.9.8)
IoError::SerializationThe bytes are not valid postcard for the expected schema (corruption or truncation after a well-formed header)
IoError::ModelStructureMismatchThe model you built does not match the file. The cause is a wrong layer count, a wrong layer type at some position, or a weight shape that disagrees with the target layer

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 you hit most often while you develop a model. RustyML raises it in 3 places: a layer-count check, a per-position type-name check, and a shape check during weight application. This third case is less obvious. A shape disagreement from a layer’s set_weights call gets wrapped into the same variant. For example, a Dense::new(2, 2, ...) file loaded into a Dense::new(3, 3, ...) target passes the count and type checks. It then fails on shape, still as ModelStructureMismatch. The message string tells you which check failed. Matching the error is straightforward:

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

fn main() {
    // Save a 1-layer model.
    let mut saved = Sequential::new();
    saved.add(Dense::new(2, 2, Linear::new()).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 = Sequential::new();
    wrong
        .add(Dense::new(2, 2, Linear::new()).unwrap())
        .add(Dense::new(2, 2, Linear::new()).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. Two 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 function discipline matters.

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 dimensions (varint-encoded), the enum discriminants (1 byte each for these small enums), and the short metadata strings. 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 one buffered write, so I/O time, not CPU time, dominates the cost.

The format is portable across machines. postcard defines its own byte order instead of dumping native-endian memory. So a checkpoint written on one architecture deserializes correctly on another, regardless of the host’s endianness. 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, 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 that epoch’s own weight update. So the last entry describes weights the model no longer holds, and a rule based on it would checkpoint a round late. evaluate runs one 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::sequential::Sequential;

fn make_arch() -> Sequential {
    let mut m = Sequential::new();
    m.add(Dense::new(4, 3, Linear::new()).unwrap())
        .add(Dense::new(3, 2, Linear::new()).unwrap());
    m
}

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, 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. So the score of the restored checkpoint must 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, so the two programs cannot disagree on the architecture. The type-name and count 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, since it only calls predict.

Manage versions across crate upgrades. Before the file header existed, a stale checkpoint could fail silently. The file header now prevents this. Every saved model opens with a magic tag and a format version, and load_from_path validates both before it decodes anything else. So a checkpoint written by a release with a different weight layout 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 weight extents. A stale file can satisfy all 3 checks by coincidence. For example, a square convolution kernel keeps its extents when its axes are permuted. A Dense weight shape stays the same no matter what tensor layout produced the input feeding it.

This guard depends on developer discipline. RustyML bumps the version whenever a weight container’s tensor layout, rank, or field order changes. A change that gets a version bump is caught. A change that someone forgot to bump is not caught, because postcard is still non-self-describing and nothing else in the file is labeled. For a checkpoint you need to reload weeks or versions later, pin the RustyML version in Cargo.toml. After a deliberate upgrade, re-run training, or load the file under the old version and re-save it 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. For more detail on the persistence machinery, see 7.2. Model Persistence in Depth. That section covers the get_weights inspection path, the SerializableSequential wrapper, and how RustyML applies weights back through downcasting.