mnist_cnn.rs

//! MNIST with two conv blocks and a linear head, around 99% test accuracy.
//!
//!     cargo run --example mnist_cnn --release [--features cuda]

use fastnn::data::{Mnist, Split};
use fastnn::prelude::*;

const EPOCHS: usize = 2;
const BATCH: usize = 64;
const TEST_BATCH: usize = 500;
const LR: f32 = 1e-3;
const SEED: u64 = 42;
const REPORT_EVERY: usize = 100;

fn main() -> fastnn::Result<()> {
    manual_seed(SEED);

    let train = Mnist::load(Split::Train)?;
    let test = Mnist::load(Split::Test)?;

    let device = Device::best();
    let model = Sequential::new()
        // 28×28 → 14×14, 1 channel → 16
        .add(Conv2d::same(1, 16, 3))
        .add(BatchNorm2d::new(16))
        .add(ReLU)
        .add(MaxPool2d::new(2))
        // 14×14 → 7×7, 16 channels → 32
        .add(Conv2d::same(16, 32, 3))
        .add(BatchNorm2d::new(32))
        .add(ReLU)
        .add(MaxPool2d::new(2))
        .add(Flatten::new())
        .add(Linear::new(32 * 7 * 7, 10));
    model.to_device(device);

    let mut opt = Adam::new(model.parameters(), LR);
    let schedule = CosineAnnealing::new(LR, EPOCHS * (train.len() / BATCH));
    let train_loader = DataLoader::new(&train, BATCH)
        .shuffle(true)
        .drop_last(true)
        .to_device(device);
    let test_loader = DataLoader::new(&test, TEST_BATCH).to_device(device);

    println!("model: {} parameters on {device}\n", model.num_parameters());

    let mut step = 0usize;
    for epoch in 1..=EPOCHS {
        model.train();
        let (mut total_loss, mut correct, mut seen) = (0.0f32, 0usize, 0usize);

        for (index, batch) in train_loader.iter().enumerate() {
            opt.set_lr(schedule.lr_at(step));
            step += 1;

            let labels = batch.labels();
            let logits = model.forward(&batch.inputs);
            let loss = cross_entropy(&logits, &labels);

            opt.zero_grad();
            loss.backward();
            opt.step();

            total_loss += loss.item();
            correct += count_correct(&logits, &labels);
            seen += batch.len();

            if (index + 1) % REPORT_EVERY == 0 {
                println!(
                    "  epoch {epoch}  batch {:4}/{}  loss {:.4}  acc {:.2}%  lr {:.2e}",
                    index + 1,
                    train_loader.batches(),
                    total_loss / (index + 1) as f32,
                    100.0 * correct as f32 / seen as f32,
                    opt.lr(),
                );
            }
        }

        // eval() makes batch norm use its running stats.
        model.eval();
        println!(
            "epoch {epoch} done  test acc {:.2}%\n",
            100.0 * accuracy(&model, &test_loader)
        );
    }

    save(&model, "mnist_cnn.fdl")?;
    println!("saved mnist_cnn.fdl");
    Ok(())
}

fn accuracy(model: &dyn Module, loader: &DataLoader) -> f32 {
    let (mut correct, mut seen) = (0usize, 0usize);
    for batch in loader.iter() {
        let logits = no_grad(|| model.forward(&batch.inputs));
        correct += count_correct(&logits, &batch.labels());
        seen += batch.len();
    }
    correct as f32 / seen as f32
}

fn count_correct(logits: &Tensor, labels: &[usize]) -> usize {
    logits
        .argmax(1)
        .iter()
        .zip(labels)
        .filter(|(p, t)| p == t)
        .count()
}