AliMokh/inklab-mnist-cnn

InkLab MNIST CNN

A small convolutional neural network trained from scratch for InkLab, an interactive drawing and robustness experiment. No pretrained weights or external inference API are used.

Results

• Official held-out MNIST test accuracy: 98.89% (10,000 images).
• Best validation accuracy: 98.74% (5,000 images).
• Training: 55,000 images, 5 epochs, seed 42.
• Parameters: 224,714; float32 safetensors checkpoint, about 0.9 MB.
• Training device: Apple MPS; measured training and evaluation time 23.24 seconds.

The 60,000 official training images were randomly split with seed 42 into 55,000 training and 5,000 validation images. The best checkpoint was selected by validation accuracy. The official 10,000-image test set was evaluated only after checkpoint selection. Full epoch losses, metrics, and confusion matrix (rows=true, columns=predicted) are in metrics.json.

Architecture and training

Three 3×3 convolutions with channels 16, 32, 64, ReLU activations, and two 2×2 max pools. A 3,136→64→10 classifier uses ReLU and dropout 0.2. AdamW, learning rate 0.001, weight decay 0.0001, batch size 256. Training augmentation uses rotations of approximately ±10 degrees, affine scales 0.9–1.1, and translations of up to 0.08 in normalized grid coordinates (about 1.1 pixels). No augmentation is applied to validation/test data. Inputs are grayscale 28×28 images, white foreground on black, float32 in [0, 1].

Reproduce training

Clone this repository, then run:

python -m pip install -r requirements.txt
python train.py --epochs 5 --data-dir /tmp/inklab-data

Uses Apple MPS if available, otherwise CPU. Seed is fixed; exact results may differ by hardware and PyTorch kernel implementation. The ten bundled images are the first example of each class in the shuffled training partition, not selected from the test set.

Inference

Download network.py and model.safetensors from the same revision, then:

import torch
from safetensors.torch import load_file
from network import InkNet

model = InkNet().eval()
model.load_state_dict(load_file("model.safetensors"))
# x: float32 tensor, shape [batch, 1, 28, 28], white ink, black background
with torch.inference_mode():
    probabilities = model(x).softmax(dim=1)

Intended use and limitations

For teaching image classification, distribution shift and model sensitivity. This model only classifies digits 0–9. It cannot reject every out-of-distribution input: letters, random scribbles and badly rotated images may receive confident digit scores. The softmax outputs are not calibrated probabilities. MNIST accuracy does not measure accuracy on browser drawings. The demo separately crops, resizes and centres canvas ink; this preprocessing is not part of the official test benchmark.

The occlusion heatmap measures positive drops in the selected class score when 4×4 patches are erased. It is an intervention-based sensitivity visualization, not a proof of causality or a complete account of the network's reasoning.

Data attribution

MNIST: Yann LeCun, Corinna Cortes and Christopher J. C. Burges. Data downloaded using torchvision's MNIST loader (checksummed public MNIST files). Dataset card and license lists MIT. Code and trained weights are released under MIT.