MNIST Digit Classification: 3 Model Approaches Compared

# MNIST Handwritten Digit Classification — 3 Approaches, Run in Parallel **Goal:** train and compare three model approaches on the classic MNIST benchmark: | # | Approach | Description | |---|----------|-------------| | 1 | **Logistic Regression** | Classical linear baseline (sklearn, lbfgs) | | 2 | **MLP** | Fully-connected net, 2 hidden layers (PyTorch) | | 3 | **CNN** | Small LeNet-style conv net (PyTorch) | **Setup:** MNIST 60k/10k → train 50k / validation 10k / test 10k; pixels normalized to [0, 1]. Each approach runs on its **own isolated branch, in parallel** (fork below), then we converge on the best and compare side-by-side.

import json, os, time, random, warnings
import numpy as np
import matplotlib.pyplot as plt
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report

import torch
import torch.nn as nn
from torch.utils.data import DataLoader, TensorDataset

warnings.filterwarnings("ignore")
torch.manual_seed(42)
np.random.seed(42)
random.seed(42)
torch.set_num_threads(min(4, os.cpu_count() or 4))  # don't thrash the sandbox CPU across parallel kernels

print("torch", torch.__version__, "| cpus:", os.cpu_count())

torch 2.7.1+cpu | cpus: 8

DATA_DIR = "/home/user/data"
RESULTS_DIR = "/home/user/results"
os.makedirs(DATA_DIR, exist_ok=True)
os.makedirs(RESULTS_DIR, exist_ok=True)

_cache = {}

def load_mnist():
    """Load MNIST -> train 50k / val 10k / test 10k, float32 in [0,1]. Cached & idempotent."""
    if "mnist" in _cache:
        return _cache["mnist"]
    try:
        from torchvision import datasets, transforms
        tr = datasets.MNIST(DATA_DIR, train=True, download=True, transform=transforms.ToTensor())
        te = datasets.MNIST(DATA_DIR, train=False, download=True, transform=transforms.ToTensor())
        X_tr, y_tr = tr.data.numpy(), tr.targets.numpy()
        X_te, y_te = te.data.numpy(), te.targets.numpy()
    except Exception as exc:
…

0%|          | 0.00/9.91M [00:00<?, ?B/s]

fig, axes = plt.subplots(2, 5, figsize=(11, 4.5))
for i, ax in enumerate(axes.flat):
    ax.imshow(X_tr[i], cmap="gray_r")
    ax.set_title(f"label: {y_tr[i]}")
    ax.axis("off")
fig.suptitle("MNIST sample digits (training set)", fontsize=13)
plt.tight_layout()
plt.show()

# ─────────────────────────────────────────────────────────────────────────────
# ONE parameterized entry point — every parallel branch calls this function.
# ─────────────────────────────────────────────────────────────────────────────

def make_loaders(X_tr, y_tr, X_val, y_val, X_te, y_te, batch_size, cnn):
    def to_t(x):
        x = torch.tensor(x, dtype=torch.float32)
        return x.unsqueeze(1) if cnn else x
    ds_tr = TensorDataset(to_t(X_tr), torch.tensor(y_tr, dtype=torch.long))
    ds_va = TensorDataset(to_t(X_val), torch.tensor(y_val, dtype=torch.long))
    ds_te = TensorDataset(to_t(X_te), torch.tensor(y_te, dtype=torch.long))
    return (DataLoader(ds_tr, batch_size=batch_size, shuffle=True),
            DataLoader(ds_va, batch_size=1024),
            DataLoader(ds_te, batch_size=1024))

def count_params(m):
    return sum(p.numel() for p in m.parameters() if p.requires_grad)

…

# ── Smoke test: validate every path BEFORE forking the 3 parallel branches ──
X_tr, y_tr, X_val, y_val, X_te, y_te = load_mnist()

# 1) MLP & CNN build + forward pass + param counts
for name, model, xb in [("mlp", MLP(), torch.zeros(4, 784)),
                        ("cnn", CNN(), torch.zeros(4, 1, 28, 28))]:
    out = model(xb)
    print(f"{name:5s} forward OK -> out{tuple(out.shape)} | params: {count_params(model):,}")

# 2) torch training loop end-to-end on a tiny synthetic slice (1 epoch)
torch.manual_seed(0)
tiny_tr = TensorDataset(torch.randn(512, 784), torch.randint(0, 10, (512,)))
tiny_va = TensorDataset(torch.randn(256, 784), torch.randint(0, 10, (256,)))
r = train_torch(MLP(), DataLoader(tiny_tr, 64, shuffle=True),
                DataLoader(tiny_va, 128), DataLoader(tiny_va, 128),
                epochs=1, lr=1e-3, label="smoke-mlp")
print("training loop OK, smoke val_acc:", round(r["val_acc"], 4))

…

mlp   forward OK -> out(4, 10) | params: 235,146

# ── Comparison of the 3 parallel branches (metrics persisted by each branch) ──
import json, os
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd

results = {}
for a in ["logreg", "mlp", "cnn"]:
    with open(os.path.join(RESULTS_DIR, f"metrics_{a}.json")) as f:
        results[a] = json.load(f)

df = pd.DataFrame([
    {"approach": a,
     "test_acc": m["test_acc"], "val_acc": m["val_acc"],
     "train_time_s": m["train_time_s"], "n_params": m["n_params"],
     "epochs": m["epochs"]}
    for a, m in results.items()
]).set_index("approach").sort_values("test_acc", ascending=False)
…

──────────────────────────────────────────────────────────
          test_acc  val_acc  train_time_s  n_params  epochs
approach                                                   
cnn         0.9904   0.9892         190.4    421642     6.0
mlp         0.9761   0.9770          78.6    235146     6.0
logreg      0.9250   0.9278         104.3      7850     NaN
──────────────────────────────────────────────────────────

## Verdict: the CNN wins | Approach | Test acc | Val acc | Train time (CPU) | Params | |---|---|---|---|---| | **CNN (LeNet-style)** 🏆 | **99.04%** | 98.92% | 190 s | 421,642 | | MLP (784-256-128-10) | 97.61% | 97.70% | 79 s | 235,146 | | Logistic Regression | 92.50% | 92.78% | 104 s | 7,850 | **Key takeaways** - The **CNN** is the only approach above 99% test accuracy — it beats the MLP by **+1.43 pp** and Logistic Regression by **+6.54 pp**. - **Logistic Regression** is a remarkably strong baseline: 92.5% with only 7,850 parameters (~30× fewer than the CNN). - The **MLP** is the fastest to train (79 s) and lands within 1.4 pp of the CNN — the best accuracy/speed trade-off. - Per-class recall (charts above): the CNN is ≥98% on every digit; LogReg struggles most with the curvy digits (5, 8). All three approaches ran **in parallel on isolated branches** (fork above). The `cnn-lenet` branch is the **winner and converges here**; the other two branches are kept for side-by-side inspection.