Build a Continuous Training Pipeline in Rust

Train, version, and promote ML models using a champion/challenger pipeline in Rust.

Introduction

30 Second Summary

Machine learning models don't stay accurate forever. Data changes, patterns shift, and yesterday's predictions become today's mistakes.

In this project, you will build a continuous training pipeline in Rust that trains a digit classifier, saves it to a local model registry, and runs a champion/challenger retraining loop to decide whether a new model deserves to be promoted.

What You'll Build

A Rust CLI tool that implements an MLOps-style continuous training pipeline for an MNIST digit classifier using the Burn deep learning framework.

By the end of this project, you'll have:

  • An MLP neural network trained on the MNIST dataset with a custom training loop that logs loss and accuracy per epoch.
  • A local model registry that versions trained models and tracks their performance in a JSON manifest.
  • A champion/challenger pipeline that retrains a new model with different hyperparameters and only promotes it when it outperforms the current champion.
  • A data drift simulation that corrupts test images and triggers automatic retraining when accuracy drops below a threshold (Secret Mission).

Are there any prerequisites?

You should be comfortable with Rust fundamentals like ownership, traits, and generics. No machine learning experience is required. This project teaches every ML concept from scratch.

Set Up the Project and Dependencies

Training a machine learning model once is not enough for production systems. Data changes over time, causing model accuracy to degrade. Production ML systems need pipelines that retrain models, compare new versions against the current best, and promote improvements automatically.

In this project, you'll build that entire pipeline in Rust using Burn, a deep learning framework. This first step gets the project scaffolded with the right dependencies and file structure.

In this step, get ready to:
  • Create a new Cargo project and configure Burn dependencies.
  • Verify the setup compiles with a simple tensor operation.
  • Set up the module files for the pipeline's components.
Create the Cargo project and add dependencies

Cargo is Rust's build tool and package manager. You'll use it to create the project scaffold and pull in the Burn framework as a dependency.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type Terminal and press Enter to open it.
  • Move to your Desktop by running this command:
cd ~/Desktop
  • Create a new Cargo project and move into it by running these commands:
cargo new ct-pipeline
cd ct-pipeline

You should see Created binary (application) `ct-pipeline` package in the output.

Seeing command not found?

Make sure Rust is installed. Run rustc --version to check. If it's missing, install it from rustup.rs.

Help me install Rust and Cargo on my machine.

  • Press Cmd+Space (macOS) or the Windows key (Windows) to open your search bar.
  • Type VS Code and press Enter to open it.
  • In VS Code, click File in the top menu bar.
  • Click Open Folder.
  • Select the ct-pipeline folder on your Desktop, then click Open.

Cargo generated a Cargo.toml file with placeholder metadata and a src/main.rs with a Hello World program. You'll replace the default dependencies with the ones this project needs.

  • In the Explorer sidebar on the left, click Cargo.toml to open it.
  • Replace everything in the file with this configuration:
[package]
name = "ct-pipeline"
version = "0.1.0"
edition = "2021"

[dependencies]
# Burn deep learning framework with CPU backend and training utilities
burn = { version = "~0.20", features = ["std", "train", "vision", "ndarray"], default-features = false }
# Serialization for model configs and registry
serde = { version = "1", features = ["derive"] }
serde_json = "1"

What are these dependencies?

  • burn is the deep learning framework. The ndarray feature enables a pure-Rust CPU backend. The train feature provides training utilities like optimizers and loss functions. The vision feature includes built-in datasets like MNIST.
  • serde and serde_json handle serialization. You'll use them to save and load the model registry as JSON.
  • default-features = false keeps the build lean by only pulling in the features you explicitly listed.

Why NdArray instead of a GPU backend?

Burn supports GPU backends like WGPU, but NdArray is pure Rust with zero external dependencies. It works on any OS without GPU drivers or setup. This project focuses on the MLOps pipeline pattern, not raw training speed, so NdArray keeps the focus on the concepts.

  • Save Cargo.toml.

Awesome, I've got everything!

Make sure you've saved Cargo.toml before moving on.

I'd like to double check the full code

[package]
name = "ct-pipeline"
version = "0.1.0"
edition = "2021"

[dependencies]
# Burn deep learning framework with CPU backend and training utilities
burn = { version = "~0.20", features = ["std", "train", "vision", "ndarray"], default-features = false }
# Serialization for model configs and registry
serde = { version = "1", features = ["derive"] }
serde_json = "1"
Verify Burn compiles with a test tensor

Before building the full pipeline, you'll confirm that Burn is installed correctly by creating a simple tensor and printing it. This code is temporary and gets replaced in the next substep.

  • In the Explorer sidebar, expand the src folder and click main.rs to open it.
  • Replace everything in the file with this code:
use burn::backend::NdArray;
use burn::prelude::*;

fn main() {
    // Set up the NdArray backend for CPU computation
    type MyBackend = NdArray<f32>;
    let device = burn::backend::ndarray::NdArrayDevice::Cpu;

    // Create a simple 2x3 tensor to verify Burn works
    let tensor = Tensor::<MyBackend, 2>::from_floats(
        [[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]],
        &device,
    );

    println!("Tensor: {}", tensor);
    println!("Burn is working!");
}

What does this code do?

  • NdArray<f32> is Burn's pure-Rust CPU backend using 32-bit floats. You assign it to a type alias called MyBackend for convenience.
  • NdArrayDevice::Cpu tells Burn to run all tensor operations on the CPU.
  • Tensor::<MyBackend, 2>::from_floats creates a 2-dimensional tensor from a nested array of floats. If this prints successfully, Burn is wired up correctly.
  • Save main.rs.

Pro tip: first compile is slow

Burn has many dependencies, so the first cargo run pulls and compiles them all. This can take 3-5 minutes. Subsequent builds are much faster.

  • Build and run the project by running this command in your terminal:
cargo run

You should see your tensor values printed as a 2x3 grid, followed by Burn is working!.

Seeing a compilation error?

  • Double check that Cargo.toml has the exact dependency lines from the previous substep. A missing feature like ndarray causes import errors.
  • Make sure you saved both Cargo.toml and main.rs before running the command.
  • Check that the Burn version is ~0.20 with the tilde. This pins to the 0.20.x range and avoids pulling a breaking pre-release.

Help me debug a Burn compilation error in my Rust project.

Set up the module file structure

The pipeline has several components: a neural network model, a data loader, training logic, a model registry, and the pipeline orchestrator. Each one lives in its own Rust module file.

  • Create the five module files by running this command:
touch src/model.rs src/data.rs src/training.rs src/pipeline.rs src/registry.rs
  • Confirm the files exist by running:
ls src/

You should see data.rs, main.rs, model.rs, pipeline.rs, registry.rs, and training.rs in the output.

touch command not found on Windows?

If you're using PowerShell without Git Bash, create each file manually. Right-click the src folder in VS Code's Explorer sidebar and select New File for each one.

Help me create files on Windows without the touch command.

Now you need to tell Rust about these modules by declaring them in main.rs.

  • Open src/main.rs and replace everything with this code:
// Declare project modules
mod data;
mod model;
mod pipeline;
mod registry;
mod training;

fn main() {
    println!("Project structure ready!");
}

What do these module declarations do?

Each mod line tells the Rust compiler to look for a corresponding file in the src folder. For example, mod model; makes the compiler include src/model.rs. Without these declarations, your code in those files would be invisible to the rest of the project.

  • Save main.rs.
  • Verify everything compiles by running:
cargo build

You should see Compiling ct-pipeline... followed by Finished with no errors.

Seeing unresolved module errors?

  • Make sure every file listed in the mod declarations actually exists in the src folder. Run ls src/ to double check.
  • Confirm the file names match exactly. For example, mod training; expects a file named training.rs (not Training.rs or train.rs).

Help me fix module resolution errors in my Rust project.

Awesome, I've got everything!

Make sure you've saved main.rs and that cargo build completes with no errors.

I'd like to double check the full code

// Declare project modules
mod data;
mod model;
mod pipeline;
mod registry;
mod training;

fn main() {
    println!("Project structure ready!");
}

Your project skeleton is ready with all the dependencies and module files in place. Next up, you'll define the neural network model that powers the digit classifier.

Define the MLP Model

Your Rust project compiles with Burn and the module files are ready. Now it's time to define the neural network at the heart of your continuous training pipeline.

Burn uses Rust's derive macros to turn ordinary structs into neural network modules. This gives you type-safe layers, built-in serialization, and automatic differentiation with zero boilerplate.

In this step, get ready to:
  • Define the MLP struct with Burn's Module derive macro.
  • Create a config struct that initializes the model with default hyperparameters.
  • Implement the forward pass, loss computation, and Burn's training traits.
Create the model struct and config

An MLP (Multi-Layer Perceptron) is a neural network made of stacked linear layers. Each layer transforms its input, and activation functions like ReLU add non-linearity between them.

Dropout randomly disables neurons during training to prevent the model from memorizing the data. Together, these components form a classifier that can learn to recognize handwritten digits.

  • Open model.rs in the src/ folder of your ct-pipeline project.
  • Add the imports and model struct by copying this code into the file:
use burn::{
    nn::{Dropout, DropoutConfig, Linear, LinearConfig, Relu},
    prelude::*,
    tensor::backend::AutodiffBackend,
    train::{ClassificationOutput, InferenceStep, TrainOutput, TrainStep},
};

// MLP neural network with two hidden layers
#[derive(Module, Debug)]
pub struct Mlp<B: Backend> {
    linear1: Linear<B>,   // Input layer: 784 -> hidden_size
    linear2: Linear<B>,   // Hidden layer: hidden_size -> hidden_size
    linear3: Linear<B>,   // Output layer: hidden_size -> num_classes
    dropout: Dropout,      // Randomly disables neurons during training
    activation: Relu,      // Non-linear activation between layers
}

What does this code do?

  • #[derive(Module, Debug)] tells Burn to treat this struct as a neural network module. Burn automatically handles parameter tracking, serialization, and device placement.
  • The struct is generic over B: Backend, so the same model definition works with any Burn backend (NdArray, WGPU, etc.).
  • Three Linear layers form the core of the network. Data flows through linear1 and linear2 (hidden layers) before reaching linear3 (output layer).

The struct defines what layers the model has. Now you need a config to set the model's hyperparameters and initialize those layers.

  • Below the Mlp struct, add the config and initialization code:
#[derive(Config, Debug)]
pub struct MlpConfig {
    #[config(default = 256)]
    hidden_size: usize,
    #[config(default = 10)]
    num_classes: usize,
    #[config(default = 0.3)]
    dropout: f64,
}

impl MlpConfig {
    // Initialize an Mlp model from this config on the given device
    pub fn init<B: Backend>(&self, device: &B::Device) -> Mlp<B> {
        Mlp {
            // 28*28 = 784 input pixels from a flattened MNIST image
            linear1: LinearConfig::new(28 * 28, self.hidden_size).init(device),
            linear2: LinearConfig::new(self.hidden_size, self.hidden_size).init(device),
            linear3: LinearConfig::new(self.hidden_size, self.num_classes).init(device),
            dropout: DropoutConfig::new(self.dropout).init(),
            activation: Relu::new(),
        }
    }
}

What does the config do?

  • #[derive(Config)] auto-generates a new() constructor and makes the struct serializable. You'll use this later to save and reload model configurations.
  • The #[config(default = ...)] attributes set sensible defaults: 256 neurons per hidden layer, 10 output classes (one per digit), and 30% dropout.
  • The init method creates an Mlp by turning each config value into a real layer. LinearConfig::new(784, 256).init(device) creates a linear layer with 784 inputs and 256 outputs, allocated on the specified device.
  • Save model.rs.
  • Verify the struct and config compile by running:
cargo check

You should see a successful check with some warnings about unused imports. Those imports will be used in the next substeps.

Seeing a compilation error?

  • Check that your use burn block includes all the listed items. Missing a comma or a closing brace in the import will cause a syntax error.
  • Make sure your Cargo.toml has the train feature enabled for Burn. Without it, types like ClassificationOutput and TrainStep won't be available.
  • Why won't my model.rs compile?
Implement the forward method

The forward pass is the core of any neural network. It defines how input data flows through the model's layers to produce a prediction.

For your MLP, each 28x28 image gets flattened into a 784-element vector. That vector passes through two hidden layers (each followed by ReLU activation and dropout) before a final output layer produces 10 scores, one per digit.

  • Below the MlpConfig impl block, add the forward method inside a new impl block for Mlp:
impl<B: Backend> Mlp<B> {
    pub fn forward(&self, images: Tensor<B, 3>) -> Tensor<B, 2> {
        let [batch_size, height, width] = images.dims();
        // Flatten each 28x28 image into a 784-element vector
        let x = images.reshape([batch_size, height * width]);

        // First hidden layer: linear transform -> activation -> dropout
        let x = self.linear1.forward(x);
        let x = self.activation.forward(x);
        let x = self.dropout.forward(x);

        // Second hidden layer: same pattern
        let x = self.linear2.forward(x);
        let x = self.activation.forward(x);
        let x = self.dropout.forward(x);

        // Output layer returns raw scores (logits) for each digit class
        self.linear3.forward(x)
    }
}

What does the forward method do?

  • The input is a 3D tensor with shape [batch_size, 28, 28]. The method flattens each image into a 1D vector of 784 values before feeding it to the first linear layer.
  • Each hidden layer applies three operations in sequence: a linear transformation (matrix multiply + bias), a ReLU activation (zeroes out negatives), and dropout (randomly zeroes neurons during training).
  • The output is a 2D tensor of shape [batch_size, 10]. Each row contains 10 raw scores (logits). The highest score indicates the model's predicted digit.
  • Save model.rs and run cargo check again.

You should still see a clean check with unused-import warnings. The compiler now knows how your model transforms input images into digit predictions.

Seeing a type mismatch error?

  • Make sure the return type is Tensor<B, 2> (2D), not Tensor<B, 3>. The reshape flattens the 3D input into 2D.
  • Check that images.reshape([batch_size, height * width]) uses the destructured variables from images.dims(). A typo in the variable names will cause a compile error.
  • Help me fix a type error in my Burn forward method.
Add classification and training traits

The forward method produces raw scores, but training needs two more things. First, a way to measure how wrong the model is (a loss function). Second, trait implementations that tell Burn how to run training and validation batches.

  • In model.rs, find the closing } of the impl<B: Backend> Mlp<B> block (the one containing the forward method).
  • Add the forward_classification method inside that impl block, right after the forward method and before the closing brace:
    pub fn forward_classification(
        &self,
        images: Tensor<B, 3>,
        targets: Tensor<B, 1, Int>,
    ) -> ClassificationOutput<B> {
        // Run the forward pass to get predictions
        let output = self.forward(images);
        // Compute cross-entropy loss between predictions and true labels
        let loss = burn::nn::loss::CrossEntropyLossConfig::new()
            .init(&output.device())
            .forward(output.clone(), targets.clone());

        ClassificationOutput::new(loss, output, targets)
    }

What is cross-entropy loss?

Cross-entropy loss measures how far the model's predicted probabilities are from the correct answer. A perfect prediction gives a loss near zero. A confident wrong prediction gives a high loss.

ClassificationOutput bundles the loss, the raw output logits, and the target labels together. Burn's training system uses this bundle to compute gradients and update the model's weights.

  • Save model.rs and run cargo check to confirm the method compiles.

You should see a successful check. The final piece is wiring up Burn's training traits so the framework knows how to run backpropagation and inference on your model.

  • At the top of model.rs, add this import below your existing use burn block:
use crate::data::MnistBatch;
  • At the bottom of model.rs (after the closing } of the impl<B: Backend> Mlp<B> block), add the TrainStep and InferenceStep trait implementations:
impl<B: AutodiffBackend> TrainStep for Mlp<B> {
    type Input = MnistBatch<B>;
    type Output = ClassificationOutput<B>;

    fn step(&self, batch: MnistBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
        // Run classification, then compute gradients via backpropagation
        let item = self.forward_classification(batch.images, batch.targets);
        TrainOutput::new(self, item.loss.backward(), item)
    }
}

impl<B: Backend> InferenceStep for Mlp<B> {
    type Input = MnistBatch<B>;
    type Output = ClassificationOutput<B>;

    fn step(&self, batch: MnistBatch<B>) -> ClassificationOutput<B> {
        // Run classification without computing gradients
        self.forward_classification(batch.images, batch.targets)
    }
}

What are TrainStep and InferenceStep?

  • TrainStep tells Burn how to process a single training batch. It runs the forward pass, computes the loss, and calls .backward() on the loss to compute gradients. Those gradients tell the optimizer how to adjust the model's weights.
  • InferenceStep tells Burn how to process a validation or test batch. It runs the same forward pass but skips gradient computation, which makes it faster.
  • Notice that TrainStep requires AutodiffBackend (a backend that supports automatic differentiation), while InferenceStep only needs a regular Backend.
  • Save model.rs.

You won't see a clean compile yet. The use crate::data::MnistBatch import references the data module, which you'll build in the next step.

Seeing an error about MnistBatch not found?

That's expected. The MnistBatch struct doesn't exist yet. You'll create it in data.rs at the start of the next step. Once data.rs is complete, the entire model file will compile cleanly.

Double check that your complete model.rs matches the reference below.

Awesome, I've got everything!

Great. Double check you've saved your file before moving on.

I'd like to double check the full code

use burn::{
    nn::{Dropout, DropoutConfig, Linear, LinearConfig, Relu},
    prelude::*,
    tensor::backend::AutodiffBackend,
    train::{ClassificationOutput, InferenceStep, TrainOutput, TrainStep},
};

use crate::data::MnistBatch;

#[derive(Module, Debug)]
pub struct Mlp<B: Backend> {
    linear1: Linear<B>,
    linear2: Linear<B>,
    linear3: Linear<B>,
    dropout: Dropout,
    activation: Relu,
}

#[derive(Config, Debug)]
pub struct MlpConfig {
    #[config(default = 256)]
    hidden_size: usize,
    #[config(default = 10)]
    num_classes: usize,
    #[config(default = 0.3)]
    dropout: f64,
}

impl MlpConfig {
    pub fn init<B: Backend>(&self, device: &B::Device) -> Mlp<B> {
        Mlp {
            linear1: LinearConfig::new(28 * 28, self.hidden_size).init(device),
            linear2: LinearConfig::new(self.hidden_size, self.hidden_size).init(device),
            linear3: LinearConfig::new(self.hidden_size, self.num_classes).init(device),
            dropout: DropoutConfig::new(self.dropout).init(),
            activation: Relu::new(),
        }
    }
}

impl<B: Backend> Mlp<B> {
    pub fn forward(&self, images: Tensor<B, 3>) -> Tensor<B, 2> {
        let [batch_size, height, width] = images.dims();
        let x = images.reshape([batch_size, height * width]);

        let x = self.linear1.forward(x);
        let x = self.activation.forward(x);
        let x = self.dropout.forward(x);

        let x = self.linear2.forward(x);
        let x = self.activation.forward(x);
        let x = self.dropout.forward(x);

        self.linear3.forward(x)
    }

    pub fn forward_classification(
        &self,
        images: Tensor<B, 3>,
        targets: Tensor<B, 1, Int>,
    ) -> ClassificationOutput<B> {
        let output = self.forward(images);
        let loss = burn::nn::loss::CrossEntropyLossConfig::new()
            .init(&output.device())
            .forward(output.clone(), targets.clone());

        ClassificationOutput::new(loss, output, targets)
    }
}

impl<B: AutodiffBackend> TrainStep for Mlp<B> {
    type Input = MnistBatch<B>;
    type Output = ClassificationOutput<B>;

    fn step(&self, batch: MnistBatch<B>) -> TrainOutput<ClassificationOutput<B>> {
        let item = self.forward_classification(batch.images, batch.targets);
        TrainOutput::new(self, item.loss.backward(), item)
    }
}

impl<B: Backend> InferenceStep for Mlp<B> {
    type Input = MnistBatch<B>;
    type Output = ClassificationOutput<B>;

    fn step(&self, batch: MnistBatch<B>) -> ClassificationOutput<B> {
        self.forward_classification(batch.images, batch.targets)
    }
}

Your model is fully defined with a forward pass, loss computation, and training traits. In the next step, you'll build the data pipeline and train this model on real MNIST handwritten digit images.

Build the Data Pipeline and Train the First Model

Your MLP model is defined and ready to learn. The next challenge is feeding it data and teaching it to recognize digits.

In this step, you will build a data pipeline that converts raw MNIST images into batched tensors. Then you will write a custom training loop that runs backpropagation to update the model's weights and save a trained model to disk.

In this step, get ready to:
  • Create a batcher that normalizes and batches MNIST images into tensors.
  • Write a training loop with gradient-based optimization using the Adam optimizer.
  • Run the training, evaluate accuracy, and save the trained model with CompactRecorder.
Create the MNIST batcher

The MLP model you built in the previous step expects batched tensors as input. Raw MNIST data arrives as individual items, each containing a 28x28 pixel array and a label. A batcher bridges this gap by collecting individual items, converting them to tensors, and normalizing the pixel values.

  • Create a new file src/data.rs and add this code:
use burn::{
    data::{dataloader::batcher::Batcher, dataset::vision::MnistItem},
    prelude::*,
};

// Stateless struct that implements the batching logic
#[derive(Clone, Default)]
pub struct MnistBatcher {}

// Holds one batch of images and their labels as tensors
#[derive(Clone, Debug)]
pub struct MnistBatch<B: Backend> {
    pub images: Tensor<B, 3>,
    pub targets: Tensor<B, 1, Int>,
}

impl<B: Backend> Batcher<B, MnistItem, MnistBatch<B>> for MnistBatcher {
    fn batch(&self, items: Vec<MnistItem>, device: &B::Device) -> MnistBatch<B> {
        // Convert each image to a tensor, reshape to [1, 28, 28], and normalize
        let images = items
            .iter()
            .map(|item| TensorData::from(item.image).convert::<B::FloatElem>())
            .map(|data| Tensor::<B, 2>::from_data(data, device))
            .map(|tensor| tensor.reshape([1, 28, 28]))
            .map(|tensor| ((tensor / 255) - 0.1307) / 0.3081)
            .collect();

        // Convert each label to a 1D integer tensor
        let targets = items
            .iter()
            .map(|item| {
                Tensor::<B, 1, Int>::from_data(
                    [(item.label as i64).elem::<B::IntElem>()],
                    device,
                )
            })
            .collect();

        // Concatenate individual tensors into a single batch
        let images = Tensor::cat(images, 0);
        let targets = Tensor::cat(targets, 0);

        MnistBatch { images, targets }
    }
}

What does this code do?

  • MnistBatcher is an empty struct that implements Burn's Batcher trait. It defines how to convert a vector of MnistItem values into a single MnistBatch.
  • Each image is converted from a raw [[f32; 28]; 28] array into a Burn tensor, reshaped to [1, 28, 28], then normalized using the standard MNIST mean (0.1307) and standard deviation (0.3081).
  • Labels are converted to integer tensors. Tensor::cat concatenates all individual tensors into a single batch along dimension 0.
  • Save src/data.rs.

You will not see a change yet. The training loop in the next substep uses this batcher to feed data to the model.

Awesome, I've got everything!

Double check you have saved src/data.rs.

I'd like to double check the full code

use burn::{
    data::{dataloader::batcher::Batcher, dataset::vision::MnistItem},
    prelude::*,
};

#[derive(Clone, Default)]
pub struct MnistBatcher {}

#[derive(Clone, Debug)]
pub struct MnistBatch<B: Backend> {
    pub images: Tensor<B, 3>,
    pub targets: Tensor<B, 1, Int>,
}

impl<B: Backend> Batcher<B, MnistItem, MnistBatch<B>> for MnistBatcher {
    fn batch(&self, items: Vec<MnistItem>, device: &B::Device) -> MnistBatch<B> {
        let images = items
            .iter()
            .map(|item| TensorData::from(item.image).convert::<B::FloatElem>())
            .map(|data| Tensor::<B, 2>::from_data(data, device))
            .map(|tensor| tensor.reshape([1, 28, 28]))
            .map(|tensor| ((tensor / 255) - 0.1307) / 0.3081)
            .collect();

        let targets = items
            .iter()
            .map(|item| {
                Tensor::<B, 1, Int>::from_data(
                    [(item.label as i64).elem::<B::IntElem>()],
                    device,
                )
            })
            .collect();

        let images = Tensor::cat(images, 0);
        let targets = Tensor::cat(targets, 0);

        MnistBatch { images, targets }
    }
}
Write the custom training loop

The training loop is the engine of any ML pipeline. Each iteration performs a forward pass to compute predictions, calculates the cross-entropy loss, then runs backpropagation to compute gradients. The Adam optimizer uses those gradients to update the model's weights.

  • Create a new file src/training.rs and add this code:
use crate::data::{MnistBatch, MnistBatcher};
use crate::model::{Mlp, MlpConfig};
use burn::{
    data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset},
    nn::loss::CrossEntropyLoss,
    optim::{AdamConfig, GradientsParams},
    prelude::*,
    record::CompactRecorder,
    tensor::backend::AutodiffBackend,
};

// Entry point with a default learning rate of 1e-3
pub fn train<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
) -> (Mlp<B::InnerBackend>, f64) {
    train_with_lr::<B>(artifact_dir, device, 1e-3)
}

// Core training function with a configurable learning rate
pub fn train_with_lr<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
    learning_rate: f64,
) -> (Mlp<B::InnerBackend>, f64) {
    std::fs::create_dir_all(artifact_dir).ok();

    // Initialize model and optimizer
    let config = MlpConfig::new();
    let mut model = config.init::<B>(device);
    let mut optim = AdamConfig::new().init();

    // Build the training data loader with batching and shuffling
    let batcher = MnistBatcher::default();
    let dataloader_train = DataLoaderBuilder::new(batcher.clone())
        .batch_size(64)
        .shuffle(42)
        .num_workers(1)
        .build(MnistDataset::train());

    let num_epochs = 2;

    for epoch in 1..=num_epochs {
        let mut total_loss = 0.0;
        let mut num_batches = 0;

        for batch in dataloader_train.iter() {
            // Forward pass: compute predictions and loss
            let output = model.forward(batch.images.clone());
            let loss = CrossEntropyLoss::new(None, &output.device())
                .forward(output, batch.targets.clone());

            // Backward pass: compute gradients and update weights
            let grads = loss.clone().backward();
            let grads = GradientsParams::from_grads(grads, &model);
            model = optim.step(learning_rate, model, grads);

            total_loss += loss.into_scalar().elem::<f64>();
            num_batches += 1;
        }

        let avg_loss = total_loss / num_batches as f64;
        println!("  Epoch {}/{}: avg loss = {:.4}", epoch, num_epochs, avg_loss);
    }

    // Switch from training mode to inference mode
    let model_valid = model.valid();

    // Evaluate on test set
    let accuracy = evaluate::<B::InnerBackend>(&model_valid, device);
    println!("  Training complete. Test accuracy: {:.2}%", accuracy);

    // Save model weights as compressed MessagePack
    let model_path = format!("{}/model", artifact_dir);
    model_valid
        .save_file(&model_path, &CompactRecorder::new())
        .expect("Failed to save model");

    // Save the model config for later reconstruction
    config
        .save(format!("{}/config.json", artifact_dir))
        .expect("Failed to save config");

    (model_valid, accuracy)
}

// Run the model against the full test set and return accuracy as a percentage
pub fn evaluate<B: Backend>(model: &Mlp<B>, device: &B::Device) -> f64 {
    let batcher = MnistBatcher::default();
    let dataloader_test = DataLoaderBuilder::new(batcher)
        .batch_size(64)
        .num_workers(1)
        .build(MnistDataset::test());

    let mut correct = 0;
    let mut total = 0;

    for batch in dataloader_test.iter() {
        let output = model.forward(batch.images);
        // Get the index of the highest-scoring class for each image
        let predictions = output.argmax(1).flatten::<1>(0, 1);
        let targets = batch.targets;

        let num_items = targets.dims()[0];
        for i in 0..num_items {
            let pred = predictions.clone().slice([i..i + 1]).into_scalar();
            let target = targets.clone().slice([i..i + 1]).into_scalar();
            if pred == target {
                correct += 1;
            }
            total += 1;
        }
    }

    (correct as f64 / total as f64) * 100.0
}

What does this code do?

  • train is a thin wrapper that calls train_with_lr with a default learning rate of 1e-3. Having a separate function lets you experiment with different learning rates later in the champion/challenger pipeline.
  • The training loop iterates over 2 epochs. Each epoch loops through every batch from the data loader. For each batch, it runs the forward pass to get predictions, computes the cross-entropy loss, then calls .backward() to compute gradients.
  • model.valid() switches the model from training mode (with autodiff) to inference mode (plain backend). This is required for evaluation and saving.
  • The model is saved using CompactRecorder, which serializes weights as compressed MessagePack (.mpk.gz files). The model config is saved separately as JSON so you can reconstruct the model architecture later.
  • evaluate runs the model against the full MNIST test set. It uses argmax(1) to pick the highest-scoring class for each image, then compares predictions against targets to calculate accuracy.
  • Save src/training.rs.

You will not see a change yet. The next substep wires this into main.rs so you can run it.

Awesome, I've got everything!

Double check you have saved src/training.rs.

I'd like to double check the full code

use crate::data::{MnistBatch, MnistBatcher};
use crate::model::{Mlp, MlpConfig};
use burn::{
    data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset},
    nn::loss::CrossEntropyLoss,
    optim::{AdamConfig, GradientsParams},
    prelude::*,
    record::CompactRecorder,
    tensor::backend::AutodiffBackend,
};

pub fn train<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
) -> (Mlp<B::InnerBackend>, f64) {
    train_with_lr::<B>(artifact_dir, device, 1e-3)
}

pub fn train_with_lr<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
    learning_rate: f64,
) -> (Mlp<B::InnerBackend>, f64) {
    std::fs::create_dir_all(artifact_dir).ok();

    let config = MlpConfig::new();
    let mut model = config.init::<B>(device);
    let mut optim = AdamConfig::new().init();

    let batcher = MnistBatcher::default();
    let dataloader_train = DataLoaderBuilder::new(batcher.clone())
        .batch_size(64)
        .shuffle(42)
        .num_workers(1)
        .build(MnistDataset::train());

    let num_epochs = 2;

    for epoch in 1..=num_epochs {
        let mut total_loss = 0.0;
        let mut num_batches = 0;

        for batch in dataloader_train.iter() {
            let output = model.forward(batch.images.clone());
            let loss = CrossEntropyLoss::new(None, &output.device())
                .forward(output, batch.targets.clone());

            let grads = loss.clone().backward();
            let grads = GradientsParams::from_grads(grads, &model);
            model = optim.step(learning_rate, model, grads);

            total_loss += loss.into_scalar().elem::<f64>();
            num_batches += 1;
        }

        let avg_loss = total_loss / num_batches as f64;
        println!("  Epoch {}/{}: avg loss = {:.4}", epoch, num_epochs, avg_loss);
    }

    let model_valid = model.valid();

    let accuracy = evaluate::<B::InnerBackend>(&model_valid, device);
    println!("  Training complete. Test accuracy: {:.2}%", accuracy);

    let model_path = format!("{}/model", artifact_dir);
    model_valid
        .save_file(&model_path, &CompactRecorder::new())
        .expect("Failed to save model");

    config
        .save(format!("{}/config.json", artifact_dir))
        .expect("Failed to save config");

    (model_valid, accuracy)
}

pub fn evaluate<B: Backend>(model: &Mlp<B>, device: &B::Device) -> f64 {
    let batcher = MnistBatcher::default();
    let dataloader_test = DataLoaderBuilder::new(batcher)
        .batch_size(64)
        .num_workers(1)
        .build(MnistDataset::test());

    let mut correct = 0;
    let mut total = 0;

    for batch in dataloader_test.iter() {
        let output = model.forward(batch.images);
        let predictions = output.argmax(1).flatten::<1>(0, 1);
        let targets = batch.targets;

        let num_items = targets.dims()[0];
        for i in 0..num_items {
            let pred = predictions.clone().slice([i..i + 1]).into_scalar();
            let target = targets.clone().slice([i..i + 1]).into_scalar();
            if pred == target {
                correct += 1;
            }
            total += 1;
        }
    }

    (correct as f64 / total as f64) * 100.0
}
Run the training and save the model

Both modules are ready. Now you need to wire them into main.rs so you can run the full pipeline.

  • Open src/main.rs and replace its contents with this code:
mod data;
mod model;
mod training;

use burn::backend::{Autodiff, NdArray};

fn main() {
    // NdArray is a pure-Rust CPU backend
    type MyBackend = NdArray<f32>;
    // Autodiff wraps it to enable gradient computation
    type MyAutodiffBackend = Autodiff<MyBackend>;

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let artifact_dir = "artifacts";

    println!("=== Continuous Training Pipeline ===\n");

    // Phase 1: Initial training
    println!("--- Phase 1: Initial Training ---");
    let (model, accuracy) =
        training::train::<MyAutodiffBackend>(artifact_dir, &device);
}

What does this code do?

  • mod data and mod training register the new modules you just created so Rust can find them.
  • training::train::<MyAutodiffBackend> calls your training function with the autodiff-enabled backend. It returns the trained model and its test accuracy.
  • You will see compiler warnings about unused model and accuracy variables. These are expected and will go away in the next step when you use them with the model registry.

What happens when MNIST downloads?

The first time you run this, Burn automatically downloads the MNIST dataset to ~/.cache/burn-dataset/mnist/. This requires an internet connection for the first run only. Subsequent runs use the cached data.

  • Save src/main.rs.
  • Run the training pipeline by running this command:
cargo run

You should see output similar to this:

=== Continuous Training Pipeline ===

--- Phase 1: Initial Training ---
  Epoch 1/2: avg loss = 0.5XXX
  Epoch 2/2: avg loss = 0.2XXX
  Training complete. Test accuracy: 9X.XX%

What should I see?

The loss should decrease from epoch 1 to epoch 2, showing that the model is learning. Test accuracy should land somewhere around 95% or higher. The exact numbers vary because of random weight initialization, but a decreasing loss trend confirms the training loop is working.

  • Confirm the saved artifacts exist by running:
ls artifacts/

You should see model.mpk.gz and config.json in the output. These are your trained model weights and architecture config.

Training not starting or seeing errors?

  • If you see unresolved import errors, make sure src/data.rs and src/training.rs are saved. Confirm that mod data and mod training appear at the top of main.rs.
  • If the MNIST download fails, check your internet connection. Burn downloads the dataset on first use.
  • If you see type mismatch errors, confirm that the MnistBatch struct in data.rs has pub on both images and targets fields.
  • Ask for help debugging your training errors.

Awesome, I've got everything!

Double check you have saved src/main.rs.

I'd like to double check the full code

mod data;
mod model;
mod training;

use burn::backend::{Autodiff, NdArray};

fn main() {
    type MyBackend = NdArray<f32>;
    type MyAutodiffBackend = Autodiff<MyBackend>;

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let artifact_dir = "artifacts";

    println!("=== Continuous Training Pipeline ===\n");

    // Phase 1: Initial training
    println!("--- Phase 1: Initial Training ---");
    let (model, accuracy) =
        training::train::<MyAutodiffBackend>(artifact_dir, &device);
}

Your first model is trained and saved. Next up, you will build a model registry that tracks model versions and their accuracy scores so you can compare them.

Build the Model Registry

In the previous step, you trained an MLP classifier on MNIST and saved it to disk. The model file sits in your artifacts folder, ready to use. But there is a problem.

Every time you retrain, the new model overwrites the old one. You lose the ability to compare versions, roll back to a previous model, or track which model is actually the best. In this step, you will build a model registry that versions every trained model and tracks which one is the current champion.

In this step, get ready to:
  • Create a registry that tracks model versions with a JSON manifest file.
  • Implement methods to register new models and look up the current champion.
  • Wire the registry into your pipeline and verify it works end to end.
Define the registry structs

The registry needs three data structures. One to describe a single model version, one for the manifest that tracks all versions, and one that ties everything together with file system operations.

  • Create a new file src/registry.rs and add the struct definitions by pasting this code:
use serde::{Deserialize, Serialize};
use std::fs;

// Each model version records its accuracy and a description
#[derive(Serialize, Deserialize, Debug)]
pub struct ModelVersion {
    pub version: u32,
    pub accuracy: f64,
    pub description: String,
}

// The manifest tracks all versions and which one is the champion
#[derive(Serialize, Deserialize, Debug)]
pub struct RegistryManifest {
    pub champion_version: u32,
    pub versions: Vec<ModelVersion>,
}

// The registry manages models on disk and keeps the manifest in sync
pub struct ModelRegistry {
    artifact_dir: String,
    manifest: RegistryManifest,
}

What do these structs do?

  • ModelVersion stores the metadata for a single trained model: its version number, test accuracy, and a short description.
  • RegistryManifest is the full JSON manifest that gets persisted to disk. It knows which version is the champion and holds the list of all versions ever registered.
  • ModelRegistry is the main struct your pipeline interacts with. It holds the artifact directory path and the in-memory manifest.

The Serialize and Deserialize derives come from serde. They let you convert the structs to and from JSON automatically.

Now the registry needs a constructor that loads an existing manifest from disk or creates a fresh one.

  • Below the ModelRegistry struct, add the impl block with the new constructor by pasting this code:
impl ModelRegistry {
    pub fn new(artifact_dir: &str) -> Self {
        let manifest_path = format!("{}/registry.json", artifact_dir);

        // Load existing manifest or create a fresh one
        let manifest = if let Ok(data) = fs::read_to_string(&manifest_path) {
            serde_json::from_str(&data).unwrap_or(RegistryManifest {
                champion_version: 0,
                versions: vec![],
            })
        } else {
            RegistryManifest {
                champion_version: 0,
                versions: vec![],
            }
        };

        ModelRegistry {
            artifact_dir: artifact_dir.to_string(),
            manifest,
        }
    }
}

What does this constructor do?

  • It tries to read registry.json from the artifact directory. If the file exists and parses correctly, it loads the previous state.
  • If the file does not exist (first run), it creates an empty manifest with no versions and champion_version set to 0.
  • This means the registry is idempotent. You can restart the pipeline and it picks up where it left off.
  • Save src/registry.rs.

You will not see a change yet. The next substep adds the methods that actually register and query models.

Implement the registry methods

The registry needs to do two things: save a trained model with a version number, and tell the pipeline which model is currently the best. You will add five methods to cover registration, querying, and promotion.

  • In src/registry.rs, find the closing } of the impl ModelRegistry block. Add the register_model and save_manifest methods just above that closing brace:
    pub fn register_model(&mut self, accuracy: f64, description: &str) {
        let version = self.manifest.versions.len() as u32 + 1;

        // Copy model files to versioned directory
        let version_dir = format!("{}/v{}", self.artifact_dir, version);
        fs::create_dir_all(&version_dir).ok();

        let src_model = format!("{}/model.mpk.gz", self.artifact_dir);
        let dst_model = format!("{}/model.mpk.gz", version_dir);
        fs::copy(&src_model, &dst_model).ok();

        let src_config = format!("{}/config.json", self.artifact_dir);
        let dst_config = format!("{}/config.json", version_dir);
        fs::copy(&src_config, &dst_config).ok();

        // Record the version in the manifest
        self.manifest.versions.push(ModelVersion {
            version,
            accuracy,
            description: description.to_string(),
        });
        self.manifest.champion_version = version;

        self.save_manifest();
    }

    fn save_manifest(&self) {
        let manifest_path = format!("{}/registry.json", self.artifact_dir);
        let json = serde_json::to_string_pretty(&self.manifest).expect("Failed to serialize");
        fs::write(manifest_path, json).expect("Failed to write manifest");
    }

What does this code do?

  • register_model creates a versioned directory (like artifacts/v1) and copies the current model files into it. This is your snapshot. The original files in artifacts/ always hold the latest model, while the versioned folders preserve history.
  • The .mpk.gz extension is what Burn's CompactRecorder produces. It is MessagePack format compressed with gzip.
  • save_manifest writes the in-memory manifest to registry.json as pretty-printed JSON. This is the single source of truth for which model version is the champion.

The pipeline also needs to query the current champion and promote a challenger when it wins. These are three small methods.

  • Still inside the impl ModelRegistry block, add these methods below save_manifest:
    pub fn current_version(&self) -> u32 {
        self.manifest.champion_version
    }

    pub fn champion_accuracy(&self) -> f64 {
        self.manifest
            .versions
            .iter()
            .find(|v| v.version == self.manifest.champion_version)
            .map(|v| v.accuracy)
            .unwrap_or(0.0)
    }

    pub fn promote_challenger(&mut self, accuracy: f64, description: &str) {
        self.register_model(accuracy, description);
        println!(
            "  Challenger promoted to champion: v{} ({:.2}%)",
            self.current_version(),
            accuracy
        );
    }

What do these methods do?

  • current_version() returns the version number of the current champion.
  • champion_accuracy() looks up the champion's accuracy from the manifest. If no champion exists yet, it returns 0.0.
  • promote_challenger() delegates to register_model and prints a promotion message. This will be called by the pipeline in the next step when a challenger beats the champion.
  • Save src/registry.rs.

You will not see a change yet. The next substep wires the registry into main.rs so you can test it end to end.

Awesome, I've got everything!

Double check you have saved src/registry.rs.

I'd like to double check the full code

use serde::{Deserialize, Serialize};
use std::fs;

// Each model version records its accuracy and a description
#[derive(Serialize, Deserialize, Debug)]
pub struct ModelVersion {
    pub version: u32,
    pub accuracy: f64,
    pub description: String,
}

// The manifest tracks all versions and which one is the champion
#[derive(Serialize, Deserialize, Debug)]
pub struct RegistryManifest {
    pub champion_version: u32,
    pub versions: Vec<ModelVersion>,
}

// The registry manages models on disk and keeps the manifest in sync
pub struct ModelRegistry {
    artifact_dir: String,
    manifest: RegistryManifest,
}

impl ModelRegistry {
    pub fn new(artifact_dir: &str) -> Self {
        let manifest_path = format!("{}/registry.json", artifact_dir);

        // Load existing manifest or create a fresh one
        let manifest = if let Ok(data) = fs::read_to_string(&manifest_path) {
            serde_json::from_str(&data).unwrap_or(RegistryManifest {
                champion_version: 0,
                versions: vec![],
            })
        } else {
            RegistryManifest {
                champion_version: 0,
                versions: vec![],
            }
        };

        ModelRegistry {
            artifact_dir: artifact_dir.to_string(),
            manifest,
        }
    }

    pub fn register_model(&mut self, accuracy: f64, description: &str) {
        let version = self.manifest.versions.len() as u32 + 1;

        // Copy model files to versioned directory
        let version_dir = format!("{}/v{}", self.artifact_dir, version);
        fs::create_dir_all(&version_dir).ok();

        let src_model = format!("{}/model.mpk.gz", self.artifact_dir);
        let dst_model = format!("{}/model.mpk.gz", version_dir);
        fs::copy(&src_model, &dst_model).ok();

        let src_config = format!("{}/config.json", self.artifact_dir);
        let dst_config = format!("{}/config.json", version_dir);
        fs::copy(&src_config, &dst_config).ok();

        // Record the version in the manifest
        self.manifest.versions.push(ModelVersion {
            version,
            accuracy,
            description: description.to_string(),
        });
        self.manifest.champion_version = version;

        self.save_manifest();
    }

    fn save_manifest(&self) {
        let manifest_path = format!("{}/registry.json", self.artifact_dir);
        let json = serde_json::to_string_pretty(&self.manifest).expect("Failed to serialize");
        fs::write(manifest_path, json).expect("Failed to write manifest");
    }

    pub fn current_version(&self) -> u32 {
        self.manifest.champion_version
    }

    pub fn champion_accuracy(&self) -> f64 {
        self.manifest
            .versions
            .iter()
            .find(|v| v.version == self.manifest.champion_version)
            .map(|v| v.accuracy)
            .unwrap_or(0.0)
    }

    pub fn promote_challenger(&mut self, accuracy: f64, description: &str) {
        self.register_model(accuracy, description);
        println!(
            "  Challenger promoted to champion: v{} ({:.2}%)",
            self.current_version(),
            accuracy
        );
    }
}
Wire the registry into the pipeline and test it

The registry module is complete. Now you need to connect it to main.rs so the trained model from the previous step gets registered as the first champion.

  • In src/main.rs, find the module declarations at the top of the file.
  • Find this line:
mod training;
  • Add mod registry; above it, so the declarations now look like this:
mod data;
mod model;
mod registry;
mod training;

Now you need to create a registry instance after training and register the model as the first champion.

  • In the main() function, find the line where training::train is called:
    let (model, accuracy) =
        training::train::<MyAutodiffBackend>(artifact_dir, &device);
  • Directly below that line, add the registry code by pasting this:
    // Register as champion
    let mut reg = registry::ModelRegistry::new(artifact_dir);
    reg.register_model(accuracy, "Initial training");
    println!(
        "Champion registered: v{} with accuracy {:.2}%\n",
        reg.current_version(),
        accuracy
    );

What does this code do?

  • ModelRegistry::new(artifact_dir) creates a new registry pointed at the artifacts directory. On the first run, there is no existing manifest, so it starts empty.
  • register_model(accuracy, "Initial training") copies the saved model into a versioned folder (artifacts/v1) and writes the manifest to registry.json.
  • The print statement confirms registration by showing the version number and accuracy.
  • Save src/main.rs.
  • Run the pipeline by executing this command:
cargo run

You should see the training output from both epochs, followed by a line like Champion registered: v1 with accuracy XX.XX%. The exact accuracy depends on training randomness, but it should be above 90%.

  • Verify the registry manifest was created by running:
cat artifacts/registry.json

You should see a JSON object with champion_version set to 1 and a single entry in the versions array showing the accuracy and description.

  • Confirm the versioned model folder exists by running:
ls artifacts/v1

You should see model.mpk.gz and config.json listed in the output.

Not seeing registry.json?

  • Make sure model.mpk.gz exists in artifacts/. The registry tries to copy this file, and it is created by the training step. If it is missing, training may have failed silently.
  • Check that mod registry; is listed with the other module declarations at the top of main.rs. A missing module declaration causes a compilation error.
  • If you see a serialization error, make sure serde_json is listed in your Cargo.toml dependencies.
  • Why isn't my model registry creating registry.json?

Awesome, I've got everything!

Double check you have saved src/main.rs.

I'd like to double check the full code

mod data;
mod model;
mod registry;
mod training;

use burn::backend::{Autodiff, NdArray};

fn main() {
    type MyBackend = NdArray<f32>;
    type MyAutodiffBackend = Autodiff<MyBackend>;

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let artifact_dir = "artifacts";

    println!("=== Continuous Training Pipeline ===\n");

    // Phase 1: Initial training
    println!("--- Phase 1: Initial Training ---");
    let (model, accuracy) =
        training::train::<MyAutodiffBackend>(artifact_dir, &device);

    // Register as champion
    let mut reg = registry::ModelRegistry::new(artifact_dir);
    reg.register_model(accuracy, "Initial training");
    println!(
        "Champion registered: v{} with accuracy {:.2}%\n",
        reg.current_version(),
        accuracy
    );
}

Your model registry is tracking versions and persisting a champion. Next up, you will build the champion/challenger pipeline that trains a new model with different hyperparameters and only promotes it if it beats the current champion.

Implement the Champion/Challenger Pipeline

Your model registry can now version and track trained models. The missing piece is the decision engine that decides when to promote a new model.

In production MLOps, you never blindly replace a deployed model. Instead, you train a new "challenger" model, compare it against the current "champion," and only promote the challenger if it actually performs better. This is the champion/challenger pattern, and you are about to build it.

In this step, get ready to:
  • Write an evaluate function that measures model accuracy on the test set.
  • Build a retraining pipeline that trains a challenger model and compares it against the champion.
  • Wire the full continuous training workflow into a CLI that prints a clear promotion decision.
Write the evaluate function

To compare two models, you need a consistent measurement. The evaluate function loads the full MNIST test set, runs every image through the model, and counts how many predictions match the actual labels.

  • In src/training.rs, add the evaluate function at the bottom of the file, after the existing train function:
pub fn evaluate<B: Backend>(model: &Mlp<B>, device: &B::Device) -> f64 {
    let batcher = MnistBatcher::default();
    let dataloader_test = DataLoaderBuilder::new(batcher)
        .batch_size(64)
        .num_workers(1)
        .build(MnistDataset::test());

    let mut correct = 0;
    let mut total = 0;

    for batch in dataloader_test.iter() {
        // Run forward pass and get predicted class indices
        let output = model.forward(batch.images);
        let predictions = output.argmax(1).flatten::<1>(0, 1);
        let targets = batch.targets;

        // Compare each prediction against the true label
        let num_items = targets.dims()[0];
        for i in 0..num_items {
            let pred = predictions.clone().slice([i..i + 1]).into_scalar();
            let target = targets.clone().slice([i..i + 1]).into_scalar();
            if pred == target {
                correct += 1;
            }
            total += 1;
        }
    }

    (correct as f64 / total as f64) * 100.0
}

What does this code do?

  • It builds a test dataloader using MnistDataset::test(), which is a separate dataset the model has never trained on.
  • For each batch, argmax(1) picks the digit class with the highest confidence score.
  • It compares every prediction against the true label and returns the percentage of correct answers.
  • Save src/training.rs.

You won't see a change yet. The evaluate function gets called when you wire up the full pipeline in Substep 3.

Build the champion/challenger pipeline

The pipeline needs to train a challenger model with different hyperparameters, then compare it against the current champion. To support this, your training module needs a way to accept a custom learning rate.

  • In src/training.rs, add this train wrapper function directly above your existing train function:
pub fn train<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
) -> (Mlp<B::InnerBackend>, f64) {
    train_with_lr::<B>(artifact_dir, device, 1e-3)
}

What does this wrapper do?

This new train function delegates to train_with_lr with a default learning rate of 1e-3. The challenger pipeline can call train_with_lr directly with a different rate to create variation.

Now refactor your existing train function (the one below the new wrapper) into train_with_lr.

  • Find the function signature of your existing train function:
pub fn train<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
) -> Mlp<B::InnerBackend> {
  • Rename it to train_with_lr, add a learning_rate: f64 parameter, and update the return type to include accuracy:
pub fn train_with_lr<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
    learning_rate: f64,
) -> (Mlp<B::InnerBackend>, f64) {
  • Inside train_with_lr, find the optimizer step line inside the batch loop:
model = optim.step(1e-3, model, grads);
  • Replace the hardcoded 1e-3 with the learning_rate parameter:
model = optim.step(learning_rate, model, grads);
  • Find the let model_valid = model.valid(); line near the bottom of the function. Add the evaluate call directly below it, before the model save code:
    let model_valid = model.valid();

    // Evaluate on test set
    let accuracy = evaluate::<B::InnerBackend>(&model_valid, device);
    println!("  Training complete. Test accuracy: {:.2}%", accuracy);

Why call evaluate inside train_with_lr?

By evaluating immediately after training, every call to train_with_lr returns both the model and its accuracy score. The pipeline can then compare scores directly without re-evaluating.

  • Find the last line of the function where it returns the model. Update it to return a tuple with the accuracy:
    model_valid
  • Replace it with:
    (model_valid, accuracy)
  • Save src/training.rs.

Awesome, I've got everything!

Double check you've saved your file before moving on.

I'd like to double check the full code

Here is the complete src/training.rs with all changes applied:

use crate::data::{MnistBatch, MnistBatcher};
use crate::model::{Mlp, MlpConfig};
use burn::{
    data::{dataloader::DataLoaderBuilder, dataset::vision::MnistDataset},
    nn::loss::CrossEntropyLoss,
    optim::{AdamConfig, GradientsParams},
    prelude::*,
    record::CompactRecorder,
    tensor::backend::AutodiffBackend,
};

pub fn train<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
) -> (Mlp<B::InnerBackend>, f64) {
    train_with_lr::<B>(artifact_dir, device, 1e-3)
}

pub fn train_with_lr<B: AutodiffBackend>(
    artifact_dir: &str,
    device: &B::Device,
    learning_rate: f64,
) -> (Mlp<B::InnerBackend>, f64) {
    std::fs::create_dir_all(artifact_dir).ok();

    let config = MlpConfig::new();
    let mut model = config.init::<B>(device);
    let mut optim = AdamConfig::new().init();

    let batcher = MnistBatcher::default();
    let dataloader_train = DataLoaderBuilder::new(batcher.clone())
        .batch_size(64)
        .shuffle(42)
        .num_workers(1)
        .build(MnistDataset::train());

    let num_epochs = 2;

    for epoch in 1..=num_epochs {
        let mut total_loss = 0.0;
        let mut num_batches = 0;

        for batch in dataloader_train.iter() {
            let output = model.forward(batch.images.clone());
            let loss = CrossEntropyLoss::new(None, &output.device())
                .forward(output, batch.targets.clone());

            let grads = loss.clone().backward();
            let grads = GradientsParams::from_grads(grads, &model);
            model = optim.step(learning_rate, model, grads);

            total_loss += loss.into_scalar().elem::<f64>();
            num_batches += 1;
        }

        let avg_loss = total_loss / num_batches as f64;
        println!("  Epoch {}/{}: avg loss = {:.4}", epoch, num_epochs, avg_loss);
    }

    let model_valid = model.valid();

    // Evaluate on test set
    let accuracy = evaluate::<B::InnerBackend>(&model_valid, device);
    println!("  Training complete. Test accuracy: {:.2}%", accuracy);

    // Save model
    let model_path = format!("{}/model", artifact_dir);
    model_valid
        .save_file(&model_path, &CompactRecorder::new())
        .expect("Failed to save model");

    config
        .save(format!("{}/config.json", artifact_dir))
        .expect("Failed to save config");

    (model_valid, accuracy)
}

pub fn evaluate<B: Backend>(model: &Mlp<B>, device: &B::Device) -> f64 {
    let batcher = MnistBatcher::default();
    let dataloader_test = DataLoaderBuilder::new(batcher)
        .batch_size(64)
        .num_workers(1)
        .build(MnistDataset::test());

    let mut correct = 0;
    let mut total = 0;

    for batch in dataloader_test.iter() {
        let output = model.forward(batch.images);
        let predictions = output.argmax(1).flatten::<1>(0, 1);
        let targets = batch.targets;

        let num_items = targets.dims()[0];
        for i in 0..num_items {
            let pred = predictions.clone().slice([i..i + 1]).into_scalar();
            let target = targets.clone().slice([i..i + 1]).into_scalar();
            if pred == target {
                correct += 1;
            }
            total += 1;
        }
    }

    (correct as f64 / total as f64) * 100.0
}

Now create the pipeline module that orchestrates the champion/challenger comparison.

  • Open src/pipeline.rs and add the retrain_and_evaluate function:
use crate::registry::ModelRegistry;
use crate::training;
use burn::tensor::backend::AutodiffBackend;

pub fn retrain_and_evaluate<B: AutodiffBackend>(
    registry: &mut ModelRegistry,
    artifact_dir: &str,
    device: &B::Device,
) {
    let champion_accuracy = registry.champion_accuracy();
    println!("  Current champion: v{} ({:.2}%)", registry.current_version(), champion_accuracy);

    // Train challenger with different learning rate
    println!("\n  Training challenger model (lr=5e-4)...");
    let challenger_dir = format!("{}/challenger", artifact_dir);
    let (_challenger_model, challenger_accuracy) =
        training::train_with_lr::<B>(&challenger_dir, device, 5e-4);

    println!("\n  --- Evaluation Results ---");
    println!("  Champion:   v{} -> {:.2}%", registry.current_version(), champion_accuracy);
    println!("  Challenger:       -> {:.2}%", challenger_accuracy);

    // Only promote if challenger outperforms champion
    if challenger_accuracy > champion_accuracy {
        println!("  Decision: PROMOTE challenger");

        // Copy challenger model to main artifact dir
        let src = format!("{}/model.mpk.gz", challenger_dir);
        let dst = format!("{}/model.mpk.gz", artifact_dir);
        std::fs::copy(&src, &dst).ok();

        let src_cfg = format!("{}/config.json", challenger_dir);
        let dst_cfg = format!("{}/config.json", artifact_dir);
        std::fs::copy(&src_cfg, &dst_cfg).ok();

        registry.promote_challenger(challenger_accuracy, "Challenger with lr=5e-4");
    } else {
        println!("  Decision: KEEP champion (challenger did not improve)");
    }
}

What does this code do?

  • It reads the champion's accuracy from the registry, so it knows what score to beat.
  • It trains a new challenger model with a learning rate of 5e-4 (half the champion's 1e-3). This creates genuine variation between the two models.
  • If the challenger scores higher, the function copies the challenger's saved files into the main artifact directory and registers it as the new champion. Otherwise, the champion stays.

Why a different learning rate?

In a real MLOps system, the challenger might differ in architecture, data, or training duration. Here, adjusting the learning rate is the simplest way to produce a meaningfully different model so the comparison logic gets exercised.

  • Save src/pipeline.rs.

You won't see a change yet. The pipeline function gets called from main.rs in the next substep.

Seeing a compile error about train_with_lr?

  • Make sure you renamed your original train function to train_with_lr in src/training.rs. The new train wrapper calls train_with_lr, so both must exist.
  • Check that train_with_lr is marked pub. The pipeline module needs to call it from a different file.
  • help me fix a compile error with train_with_lr in my training module

Awesome, I've got everything!

Double check you've saved your file before moving on.

I'd like to double check the full code

Here is the complete src/pipeline.rs:

use crate::registry::ModelRegistry;
use crate::training;
use burn::tensor::backend::AutodiffBackend;

pub fn retrain_and_evaluate<B: AutodiffBackend>(
    registry: &mut ModelRegistry,
    artifact_dir: &str,
    device: &B::Device,
) {
    let champion_accuracy = registry.champion_accuracy();
    println!("  Current champion: v{} ({:.2}%)", registry.current_version(), champion_accuracy);

    // Train challenger with different learning rate
    println!("\n  Training challenger model (lr=5e-4)...");
    let challenger_dir = format!("{}/challenger", artifact_dir);
    let (_challenger_model, challenger_accuracy) =
        training::train_with_lr::<B>(&challenger_dir, device, 5e-4);

    println!("\n  --- Evaluation Results ---");
    println!("  Champion:   v{} -> {:.2}%", registry.current_version(), champion_accuracy);
    println!("  Challenger:       -> {:.2}%", challenger_accuracy);

    // Only promote if challenger outperforms champion
    if challenger_accuracy > champion_accuracy {
        println!("  Decision: PROMOTE challenger");

        // Copy challenger model to main artifact dir
        let src = format!("{}/model.mpk.gz", challenger_dir);
        let dst = format!("{}/model.mpk.gz", artifact_dir);
        std::fs::copy(&src, &dst).ok();

        let src_cfg = format!("{}/config.json", challenger_dir);
        let dst_cfg = format!("{}/config.json", artifact_dir);
        std::fs::copy(&src_cfg, &dst_cfg).ok();

        registry.promote_challenger(challenger_accuracy, "Challenger with lr=5e-4");
    } else {
        println!("  Decision: KEEP champion (challenger did not improve)");
    }
}
Wire everything into the CLI

The evaluate function and pipeline module are ready. Now you connect them into a single CLI that runs the full continuous training workflow: train a champion, register it, then run a retraining cycle.

  • Open src/main.rs and replace its contents with the following:
mod data;
mod model;
mod pipeline;
mod registry;
mod training;

use burn::backend::{Autodiff, NdArray};

fn main() {
    type MyBackend = NdArray<f32>;
    type MyAutodiffBackend = Autodiff<MyBackend>;

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let artifact_dir = "artifacts";

    println!("=== Continuous Training Pipeline ===\n");

    // Phase 1: Initial training
    println!("--- Phase 1: Initial Training ---");
    let (model, accuracy) =
        training::train::<MyAutodiffBackend>(artifact_dir, &device);

    // Register as champion
    let mut reg = registry::ModelRegistry::new(artifact_dir);
    reg.register_model(accuracy, "Initial training");
    println!(
        "Champion registered: v{} with accuracy {:.2}%\n",
        reg.current_version(),
        accuracy
    );

    // Phase 2: Retraining cycle
    println!("--- Phase 2: Retraining Cycle ---");
    pipeline::retrain_and_evaluate::<MyAutodiffBackend>(&mut reg, artifact_dir, &device);

    println!("\n=== Pipeline Complete ===");
}

What does this code do?

  • Phase 1 trains the initial model using the default learning rate and registers it as the first champion in the registry.
  • Phase 2 calls retrain_and_evaluate, which trains a challenger with a different learning rate, compares accuracy scores, and either promotes or keeps the champion.
  • The NdArray<f32> backend runs everything on CPU in pure Rust. The Autodiff wrapper adds automatic differentiation for backpropagation during training.
  • Save src/main.rs.

Awesome, I've got everything!

Double check you've saved all files before running the pipeline.

I'd like to double check the full code

Here is the complete src/main.rs:

mod data;
mod model;
mod pipeline;
mod registry;
mod training;

use burn::backend::{Autodiff, NdArray};

fn main() {
    type MyBackend = NdArray<f32>;
    type MyAutodiffBackend = Autodiff<MyBackend>;

    let device = burn::backend::ndarray::NdArrayDevice::Cpu;
    let artifact_dir = "artifacts";

    println!("=== Continuous Training Pipeline ===\n");

    // Phase 1: Initial training
    println!("--- Phase 1: Initial Training ---");
    let (model, accuracy) =
        training::train::<MyAutodiffBackend>(artifact_dir, &device);

    // Register as champion
    let mut reg = registry::ModelRegistry::new(artifact_dir);
    reg.register_model(accuracy, "Initial training");
    println!(
        "Champion registered: v{} with accuracy {:.2}%\n",
        reg.current_version(),
        accuracy
    );

    // Phase 2: Retraining cycle
    println!("--- Phase 2: Retraining Cycle ---");
    pipeline::retrain_and_evaluate::<MyAutodiffBackend>(&mut reg, artifact_dir, &device);

    println!("\n=== Pipeline Complete ===");
}
  • Run the full pipeline by running this command:
cargo run

What should I see?

The pipeline runs in two phases. Phase 1 trains the champion model over 2 epochs, printing the average loss for each epoch, then the test accuracy. Phase 2 trains the challenger with a different learning rate and prints a side-by-side comparison.

At the end, you see a Decision: PROMOTE challenger or Decision: KEEP champion line. The exact result depends on random initialization, but both outcomes prove the pipeline logic works. Expect accuracy values in the 90-97% range for both models.

Pro tip: patience

The pipeline trains two full models back to back. On the NdArray CPU backend, this can take a few minutes total. That is expected. You are running real backpropagation on 60,000 training images, twice.

Seeing a type mismatch error?

  • If you see an error about expected (Mlp<...>, f64), found Mlp<...>, make sure you updated the return type of train_with_lr to (Mlp<B::InnerBackend>, f64) and that the function returns (model_valid, accuracy) at the end.
  • If you see unresolved import pipeline, check that mod pipeline; is declared at the top of src/main.rs.
  • help me debug a compile error in my Burn continuous training pipeline

Secret mission

Simulate Data Drift and Automatic Retraining

Your pipeline trains and promotes models, but what happens when the real world changes? In production, data drift silently degrades model performance over time. Add a drift simulation that corrupts test images with random noise, detects when accuracy drops below a threshold, and automatically triggers retraining to recover.

Before you go

Clean Up Your Resources

Decide whether to keep your resources running, pause them to come back later, or delete them entirely. This project runs entirely locally, so there are no ongoing costs.

Resources you used:

  • Cargo project directory (ct-pipeline).
  • Trained model artifacts in the artifacts/ directory (versioned models, registry manifest, challenger files).
  • Cached MNIST dataset at ~/.cache/burn-dataset/mnist/.

Keep everything running

No action needed. Choose this if you're still actively building or plan to extend the pipeline with the Secret Mission.

  • Your project files and trained models will stay in place.
  • The cached MNIST dataset will remain available for future runs without re-downloading.

Pause - I'll come back to this later

Free up disk space by removing generated artifacts while keeping your source code intact.

  • Delete the artifacts/ directory. You can retrain and regenerate all models by running the pipeline again.
  • Run cargo clean to remove compiled build artifacts from the target/ directory.
  • Keep the MNIST cache at ~/.cache/burn-dataset/mnist/ so you don't have to download it again next time.

Delete - I don't want to use this again

Remove all project resources and start fresh.

  • Delete the entire ct-pipeline project directory.
  • Delete the cached MNIST dataset at ~/.cache/burn-dataset/mnist/.
  • Your Rust toolchain and Cargo installation are shared system tools. Keep them if you plan to use Rust for other projects.

Mission accomplished

That's a Wrap

Congratulations, you just built an MLOps pipeline entirely in Rust! You created a continuous training system that trains, versions, evaluates, and selectively promotes ML models. All with zero cloud dependencies.

You've learned how to:

  • Train an MLP classifier on the MNIST dataset using Burn's NdArray backend with a custom training loop that handles forward passes, backpropagation, and gradient updates.
  • Build a local model registry that versions trained models with CompactRecorder serialization and tracks accuracy scores in a JSON manifest.
  • Implement a champion/challenger evaluation loop that trains a challenger model with different hyperparameters and only promotes it when it outperforms the current champion.
  • Simulate data drift by corrupting test images with random noise and triggering automatic retraining when accuracy drops below a threshold.

Ready to quiz yourself?