MNIST Digit Recognizer — Bangrui Fang

Model Architecture

Input 28 × 28
Grayscale → [1, 1, 28, 28]

Conv2d + ReLU 32 filters
kernel 3 × 3 → [1, 32, 26, 26]

Conv2d + ReLU 64 filters
kernel 3 × 3 → [1, 64, 24, 24]

MaxPool2d 2 × 2
Dropout 0.25 → [1, 64, 12, 12]

Flatten → Linear 9 216 → 128
ReLU · Dropout 0.5 → [1, 128]

Linear · Softmax 128 → 10
class scores → [1, 10]

Why I built this

Most ML tutorials hand you a pre-trained model and a notebook. You run the cells, get 99% accuracy, and learn almost nothing about what actually happened. I wanted to go through the full pipeline myself — from writing the training loop to serving inference in a browser — so I could understand where each decision matters and where things break.

The goal was not to build the best MNIST classifier. It was to build a complete system I can reason about — training, export, preprocessing, and runtime — with no hidden steps.

Constraints I set for myself

Train from scratch — no transfer learning, no pretrained weights
No server for inference — everything runs client-side
Match my own PyTorch preprocessing exactly in JavaScript
Model must load and run on mobile browsers
Keep the model under 5 MB for reasonable load times

Preprocessing pipeline

The preprocessing step is where most browser-based ML demos silently break. If the normalization in your browser code doesn't exactly match your training pipeline, the model sees a completely different input distribution and confidence scores become meaningless.

My PyTorch training uses Normalize((0.1307,), (0.3081,)) — the MNIST dataset's channel mean and standard deviation. In the browser, every drawing goes through the same transform: bounding box crop, scale to 20px on the longest side, center in a 28×28 canvas with 4px padding, then (pixel / 255 − 0.1307) / 0.3081 per pixel.

Getting this right required debugging with side-by-side comparisons between the Python and JavaScript preprocessing output. A common mistake is forgetting the centering step — MNIST digits are centered in their bounding boxes, and the model expects that.

4.6 MB Model file size

< 5 ms Inference latency

~24M FLOPs per pass

93K Trainable parameters

Failure cases & limitations

The model is strong on clean, centered digits but has predictable failure modes. Understanding where it breaks is more interesting than the accuracy number.

Ambiguous shapes — a poorly drawn 4 can look like a 9, and the model reflects that ambiguity with split confidence. This is correct behavior.
Very thick strokes — the model was trained on thin strokes. Drawing with a very thick line changes the effective shape and drops confidence.
Off-center digits — despite augmentation, extreme offset still confuses the model. The preprocessing crops and centers, but very small or very large drawings can end up poorly scaled.
Non-digit input — drawing letters or symbols produces high-confidence wrong answers. The model has no "none of the above" class.
Domain gap — MNIST was scanned and anti-aliased; mouse/touchscreen drawings have different stroke characteristics. Data augmentation helps but doesn't fully close this gap.

What this taught me

A model's failure modes tell you more than its accuracy. 99% on the test set sounds impressive until you realize the test set comes from the same distribution as the training data. Real-world input — messy mouse drawings — is always harder. The gap between benchmark accuracy and deployed reliability is where engineering judgment matters most.

What I'd improve next

Add a rejection threshold — if max confidence is below a threshold, show "unsure" instead of forcing a prediction. This is more honest than always producing a digit.
Collect real drawing data — fine-tune on actual mouse/touch drawings instead of relying on augmented MNIST. The domain gap matters.
Try quantization — INT8 quantization could cut the model from 4.6 MB to ~1.2 MB with minimal accuracy loss, improving mobile load times.
Add stroke normalization — normalize stroke thickness before feeding to the model, reducing sensitivity to drawing speed and input device.
WebGPU backend — ONNX Runtime Web supports WebGPU for GPU-accelerated inference. Worth benchmarking against the current WASM backend.

On scope

I intentionally kept this project small. The point was not to build the most accurate digit classifier — it was to own every step of the pipeline and know exactly where the tradeoffs are. A more complex model would improve accuracy but would obscure the learning. I'd rather ship something I fully understand.

How it runs in your browser

Training

Trained from scratch in PyTorch on the MNIST dataset

# 60,000 training images, 10,000 test images
# 20 epochs · batch size 128 · Adam lr=0.001
# data augmentation: rotation ±10°, translate 10%

train_transform = transforms.Compose([
    transforms.RandomRotation(10),
    transforms.RandomAffine(0, translate=(0.1, 0.1)),
    transforms.ToTensor(),
    transforms.Normalize((0.1307,), (0.3081,))
])

I added data augmentation after noticing the baseline model struggled with slightly off-center digits. Random rotation and translation helped generalize to messier handwriting.

Browser inference

Exported to ONNX, then run via WebAssembly — no server needed

# Export once after training
torch.onnx.export(model, dummy, "mnist.onnx")

# Browser loads the .onnx file directly
# ONNX Runtime Web compiles it to WASM at runtime
session = await ort.InferenceSession.create("mnist.onnx")

# Your drawing → 28×28 tensor → model → softmax → digit

The whole inference pipeline runs on your device. The model file is 4.6 MB — loaded once, cached by the browser. Drawing-to-prediction latency is under 5 ms on most machines.