MNIST MLP vs CNN Benchmark Comparison

import torch, torch.nn as nn, torch.nn.functional as F
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
from sklearn.metrics import confusion_matrix, classification_report
import numpy as np, json, os, time, pickle

print("torch", torch.__version__, "| cuda:", torch.cuda.is_available())
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# ---------------------------------------------------------------------------
# SHARED DATA SETUP — one identical train/test split for BOTH architectures.
# We reseed INSIDE train_model so each architecture sees the exact same
# shuffled batch order (true apples-to-apples comparison).
# ---------------------------------------------------------------------------
transform = transforms.Compose([
    transforms.ToTensor(),
    transforms.Normalize((0.5,), (0.5,))
])
…

torch 2.7.1+cpu | cuda: False

# ---------------------------------------------------------------------------
# ONE parameterized train/score function. Each branch calls this with a
# different `arch` and gets back a full result dict: loss curves, final test
# accuracy, confusion matrix, per-class metrics, per-class error counts.
# The result is ALSO persisted to a shared JSON artifact so the two branches
# can be read together for the side-by-side comparison.
# ---------------------------------------------------------------------------
def _jsonable(obj):
    if isinstance(obj, (np.integer,)): return int(obj)
    if isinstance(obj, (np.floating,)): return float(obj)
    if isinstance(obj, np.ndarray): return obj.tolist()
    if isinstance(obj, dict): return {k: _jsonable(v) for k, v in obj.items()}
    if isinstance(obj, (list, tuple)): return [_jsonable(v) for v in obj]
    return obj

RESULT_DIR = './mnist_results'
os.makedirs(RESULT_DIR, exist_ok=True)

…