MNIST CNN → SONNX C export: digit classification tutorial

This notebook trains a tiny CNN on MNIST with PyTorch, then pushes it through the aidge_export_sonnx pipeline and runs the generated C++ code on a test digit:

Requirements: aidge_core, aidge_backend_cpu, aidge_onnx, aidge_export_sonnx, a C++ toolchain with make, plus onnx, numpy, Pillow, matplotlib, torch and torchvision.

1. Setup

import dataclasses
import os
import shutil
import subprocess
import tempfile
import time
from pathlib import Path

import numpy as np
import onnx
import torch
from onnx import TensorProto, helper
from PIL import Image
from torchvision import datasets
import matplotlib.pyplot as plt

import aidge_core
import aidge_onnx
import aidge_export_sonnx

from aidge_export_sonnx.examples.mnist.graph_clm import (
    BLOCKS,
    ChainBuilder,
    compute_reference_chain,
)
from aidge_export_sonnx.examples.mnist.train_mnist_pytorch import MnistNet, train
from aidge_export_sonnx.examples.mnist.run_mnist_inference import (
    load_and_preprocess,
    parse_sonnx_results,
    softmax,
)

All generated artifacts go to a temporary folder. We chdir into it so that relative output paths used by the export API never land in the repository.

WORKDIR = Path(tempfile.mkdtemp(prefix="mnist_tutorial_"))
os.chdir(WORKDIR)
print("Working directory:", WORKDIR)

2. Train the network with PyTorch

train() downloads MNIST, trains MnistNet and exports each convolution as <block>_W.npy / <block>_B.npy (float64). We point both the dataset and the weights to the temp folder.

EPOCHS = 3
data_dir = WORKDIR / "data"
weights_dir = WORKDIR / "weights"

train(epochs=EPOCHS, data_dir=str(data_dir), weights_dir=str(weights_dir))
print(sorted(p.name for p in weights_dir.glob("*.npy")))

3. The network definition

BLOCKS is a list of ConvLeakyReluMaxPoolConfig, one per block. Each entry carries the path of the weight folder it loads from, and in the package that path points next to the source file.

blocks = [dataclasses.replace(b, weights_dir=str(weights_dir)) for b in BLOCKS]

print(
    f"{'block':12s} {'out_ch':>6s} {'kernel':>7s} {'pads':>13s} {'lrelu':>5s} {'pool':>5s}"
)
for b in blocks:
    print(
        f"{b.name:12s} {b.out_channels:6d} {str(b.kernel_hw):>7s} {str(b.conv_pads):>13s} "
        f"{'Y' if b.leaky_enabled else '-':>5s} {'Y' if b.pool_enabled else '-':>5s}"
    )

4. Pick a test digit

We take a digit from the MNIST test set (already downloaded by the training step), save it as a PNG and load it with the package’s load_and_preprocess (grayscale, 28×28, scaled to [0, 1]). MNIST digits are white on black, as the network expects. For your own photo of a dark digit on a light background, pass invert=True.

SAMPLE_INDEX = 0  # any index in [0, 9999]

test_ds = datasets.MNIST(root=str(data_dir), train=False, download=False)
pil_digit, true_label = test_ds[SAMPLE_INDEX]

image_path = WORKDIR / "digit.png"
pil_digit.save(image_path)

x = load_and_preprocess(str(image_path))
print("Network input:", x.shape, x.dtype, "| true label:", true_label)
plt.imshow(pil_digit, cmap="gray")
plt.axis("off")
plt.title(f"MNIST test #{SAMPLE_INDEX}")
plt.show()

5. Build the ONNX graph

ChainBuilder walks over the blocks, loads each .npy file and appends the Conv, LeakyRelu and MaxPool nodes. Everything is float64, which is the datatype used for the export.

input_shape = list(x.shape)
X = helper.make_tensor_value_info("X", TensorProto.DOUBLE, input_shape)

builder = ChainBuilder(input_name="X", input_shape=input_shape)
builder.add_blocks(blocks)
model = builder.build_model(X)
onnx.checker.check_model(model)

onnx_path = WORKDIR / "mnist_chain.onnx"
onnx.save(model, onnx_path)
print(f"{len(model.graph.node)} nodes:", [n.op_type for n in model.graph.node])

6. Import into aidge, compile with the SONNX backend, schedule

aidge_model = aidge_onnx.load_onnx(str(onnx_path))

aidge_model.compile(
    backend="export_sonnx",
    datatype=aidge_core.dtype.float64,
    dims=[input_shape],
)

scheduler = aidge_core.SequentialScheduler(aidge_model)
scheduler.generate_scheduling()

7. Export to C/C++, build and run

make builds bin/run_main, whose standard output contains the raw network output between SONNX_RESULTS_BEGIN / SONNX_RESULTS_END markers.

export_folder = "sonnx_export_mnist"
aidge_export_sonnx.export(
    export_folder_name=export_folder,
    scheduler=scheduler,
    inputs_tensor=[aidge_core.Tensor(x, backend="cpu")],
)

build = subprocess.run(
    ["make", "-C", export_folder, "clean", "all"], capture_output=True, text=True
)
if build.returncode != 0:
    print(build.stdout, build.stderr)
    raise RuntimeError("make failed")

binary = Path(export_folder) / "bin" / "run_main"
t0 = time.perf_counter()
run = subprocess.run(
    [str(binary.resolve())], capture_output=True, text=True, cwd=export_folder
)
elapsed = time.perf_counter() - t0
if run.returncode != 0:
    print(run.stdout, run.stderr)
    raise RuntimeError(f"run_main exited with code {run.returncode}")
print(f"Inference time (process included): {elapsed * 1000:.1f} ms")

8. Read the logits and predict

The network outputs 10 logits (one per digit). A softmax turns them into probabilities.

values = parse_sonnx_results(run.stdout)
assert values is not None, "No SONNX_RESULTS block found in the binary's output"

logits = np.array(values, dtype=np.float64).reshape(blocks[-1].out_channels)
probs = softmax(logits)
predicted = int(np.argmax(probs))
print(
    f"Predicted digit: {predicted} (confidence {probs[predicted] * 100:.1f}%) | true label: {true_label}"
)

fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(8, 3))
ax0.imshow(pil_digit, cmap="gray")
ax0.axis("off")
ax0.set_title(f"predicted: {predicted}")
ax1.bar(range(10), probs)
ax1.set_xticks(range(10))
ax1.set_xlabel("digit")
ax1.set_ylabel("probability")
plt.tight_layout()
plt.show()

9. Check the result against references

Two independent references: the NumPy implementation shipped with the package (compute_reference_chain) and the original PyTorch model, rebuilt in float64 from the same .npy files. The generated C code should match both to about 1e-6.

# (a) NumPy reference from the package
np_ref = compute_reference_chain(x, blocks).reshape(-1)
np.testing.assert_allclose(logits, np_ref, rtol=1e-6, atol=1e-6)
print("max |SONNX - NumPy|   =", float(np.abs(logits - np_ref).max()))

# (b) PyTorch, float64, same weights
net = MnistNet().double().eval()
for conv, name in [
    (net.conv1, "block1"),
    (net.conv2, "block2"),
    (net.conv3, "block3_head"),
]:
    conv.weight.data = torch.from_numpy(np.load(weights_dir / f"{name}_W.npy"))
    conv.bias.data = torch.from_numpy(np.load(weights_dir / f"{name}_B.npy"))
with torch.no_grad():
    torch_ref = net(torch.from_numpy(x))[0].numpy()
np.testing.assert_allclose(logits, torch_ref, rtol=1e-6, atol=1e-6)
print("max |SONNX - PyTorch| =", float(np.abs(logits - torch_ref).max()))
print("Outputs match.")

10. Cleanup

os.chdir(Path.home())
shutil.rmtree(WORKDIR, ignore_errors=True)
print("Temporary files removed.")