Porting a PyTorch CNN to Rust

Youtube video: https://youtu.be/JxXepbop52A?si=YAKBSOyMfB83C-l9

Most people’s first CNN lives in a Colab notebook: a few nn.Conv2d calls, .to(device), a training loop copy-pasted from a tutorial, and a plot of the loss curve at the end. That was mine too. The follow-up question I couldn’t shake was: how much of that actually depends on PyTorch’s machinery, and how much would survive a port to a language with none of it?

So I rebuilt the CNN from that original notebook — a TinyVGG-style classifier for FashionMNIST — in Rust, using burn, a deep learning framework that’s pure Rust with no libtorch dependency. The result is CNN.

The architecture, unchanged

The point of the exercise was to replicate FashionMNISTModelV2, not redesign it. Two convolutional blocks feeding a linear classifier:

conv_block_1: Conv(1→10, 3x3, pad 1) → ReLU → Conv(10→10, 3x3, pad 1) → ReLU → MaxPool(2)
conv_block_2: Conv(10→10, 3x3, pad 1) → ReLU → Conv(10→10, 3x3, pad 1) → ReLU → MaxPool(2)
classifier: Flatten → Linear(10*7*7 → 10)

With a 28x28 input, each block halves the spatial dimensions — 28 → 14 → 7 — so by the time the classifier sees the feature maps, it’s working with 10 channels of 7x7 output. Same hyperparameters as the notebook too: SGD at lr=0.1, batch size 32, 3 epochs.

Keeping the architecture and hyperparameters identical was deliberate. It meant any difference in behavior between the Rust version and the PyTorch original had to come from the framework or the port itself, not from a different model.

Where the friction actually was

PyTorch hides an enormous amount of bookkeeping behind Dataset and DataLoader. Porting that meant writing src/data.rs to handle the FashionMNIST IDX file format directly — downloading Zalando’s mirror of the dataset, parsing the binary format, and building a batcher that does the ToTensor-equivalent scaling of pixel values into [0, 1]. None of this is hard, exactly, but it’s usually invisible, and writing it by hand is the fastest way to actually understand what a DataLoader is doing for you.

The other place the port earns its keep is src/model.rs, where the model implements burn’s TrainStep and InferenceStep traits explicitly. In PyTorch, the training loop and the forward pass are yours to write in Python and it’s easy to blur what’s “the model” versus what’s “the loop.” Burn’s trait-based structure forces a cleaner separation, which — once I stopped fighting it — actually made the code easier to reason about than the notebook it came from.

Backend swapping is the actual payoff

The reason to use burn instead of just calling into libtorch from Rust is the backend abstraction. src/main.rs selects the CPU backend (NdArray, wrapped in Autodiff for gradient tracking) by default. Swapping to GPU is a one-line change to burn::backend::Wgpu plus a feature flag in Cargo.toml — no CUDA toolchain, no libtorch binary matching your platform and driver version. That’s the pitch of a pure-Rust deep learning framework: the same model code targets CPU or GPU without touching a system dependency.

Predictions with ASCII

Training gives you a live dashboard of accuracy and loss, which is table stakes at this point. The part I actually like is src/predict.rs: it loads the saved artifacts/model.mpk and artifacts/config.json, asks for a test image index, and — instead of just printing a label — draws the image as ASCII art directly in the terminal next to the prediction:

Test image index (0–9999): 0
 … (ASCII drawing of an ankle boot) …
Prediction: Ankle boot
Truth: Ankle boot
Result: correct

It’s a small thing, but rendering the actual input next to the prediction in a plain terminal — no matplotlib, no notebook cell — makes misclassifications legible in a way a bare accuracy number doesn’t. You can eyeball which ankle boots it’s getting wrong.

What this port is and isn’t

This isn’t a case for rewriting your ML pipeline in Rust for performance — burn’s CPU backend on a toy dataset isn’t the interesting part. The interesting part is that reimplementing something you already understand in a framework with fewer hidden defaults is one of the more efficient ways to find out what you didn’t actually understand the first time. If you’ve only ever trained a CNN through PyTorch’s abstractions, I’d recommend the exercise — even a small one, on a dataset as unglamorous as FashionMNIST.

Code’s on GitHub; cargo run — release gets you training or predicting from an interactive menu.