MNIST Architecture Comparison: MLP vs CNN

# Handwritten Digit Recognition — Architecture Comparison **Goal:** train two architectures on MNIST in parallel branches and pick the best one. | Variant | Architecture | Idea | |---|---|---| | **MLP** | 784 → 256 → 128 → 10 (ReLU + Dropout) | Fast fully-connected baseline, no spatial structure | | **CNN** | Conv(32) → Conv(64) → FC(128) | Learns local spatial features — the classic MNIST winner | **Protocol:** identical train/val/test split (50k / 10k / 10k), same optimizer (Adam, lr=1e-3), same epochs & batch size. Selection metric: **best validation accuracy**, tie-broken by test accuracy. Each variant runs on its own forked branch below, in parallel, saving artifacts (history, test predictions, model weights) to `/home/user/mnist_runs/<arch>/`. The analysis then converges back here for visualizations and the final verdict.

import os, time, json, random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset

SEED = 42
random.seed(SEED); np.random.seed(SEED); torch.manual_seed(SEED)
torch.set_num_threads(4)  # two variants run in parallel on 8 CPU cores
DEVICE = torch.device('cpu')
DATA_DIR = '/home/user/mnist_data'
ART_DIR  = '/home/user/mnist_runs'
os.makedirs(DATA_DIR, exist_ok=True); os.makedirs(ART_DIR, exist_ok=True)

# ---- Load MNIST (torchvision, with OpenML fallback) ----
MEAN, STD = 0.1307, 0.3081
try:
…

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

def build_model(arch: str) -> nn.Module:
    """One parameterized builder so both variants share preprocessing/training code."""
    if arch == 'mlp':
        return nn.Sequential(
            nn.Flatten(),
            nn.Linear(784, 256), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(256, 128), nn.ReLU(), nn.Dropout(0.2),
            nn.Linear(128, 10),
        )
    if arch == 'cnn':
        class CNN(nn.Module):
            def __init__(self):
                super().__init__()
                self.conv1 = nn.Conv2d(1, 32, 3, padding=1)
                self.conv2 = nn.Conv2d(32, 64, 3, padding=1)
                self.pool  = nn.MaxPool2d(2)
                self.fc1   = nn.Linear(64 * 7 * 7, 128)
                self.fc2   = nn.Linear(128, 10)
…

Scaffold ready: build_model / evaluate / train_model  |  archs: mlp, cnn

import pandas as pd

# ---- Load both variants' artifacts (saved by each forked branch) ----
runs, preds = {}, {}
for arch in ['mlp', 'cnn']:
    with open(f'{ART_DIR}/{arch}/result.json') as f:
        runs[arch] = json.load(f)
    preds[arch] = np.load(f'{ART_DIR}/{arch}/test_preds.npz')

rows = []
for arch, r in runs.items():
    s = dict(r['summary'])
    s['final_val_acc'] = r['history'][-1]['val_acc']
    rows.append(s)
cmp_df = (pd.DataFrame(rows).set_index('arch')
          [['params', 'epochs', 'train_time_s', 'best_val_acc', 'test_acc', 'test_loss']])
winner = cmp_df['best_val_acc'].idxmax()

…

🏆 Winner by validation accuracy: CNN (val 98.44%, test 98.71%)

import matplotlib.pyplot as plt

plt.rcParams.update({'figure.dpi': 110, 'font.size': 10})
COLORS = {'mlp': '#e67e22', 'cnn': '#2e86de'}

fig, axes = plt.subplots(1, 2, figsize=(12.5, 4.3))
for arch, r in runs.items():
    h = pd.DataFrame(r['history'])
    axes[0].plot(h['epoch'], h['train_loss'], 'o--', color=COLORS[arch], alpha=0.55, label=f'{arch.upper()} train')
    axes[0].plot(h['epoch'], h['val_loss'],   'o-',  color=COLORS[arch], label=f'{arch.upper()} val')
    axes[1].plot(h['epoch'], h['val_acc'] * 100, 'o-', color=COLORS[arch], lw=2, label=arch.upper())

axes[0].set_title('Loss per epoch'); axes[0].set_xlabel('Epoch'); axes[0].set_ylabel('Cross-entropy')
axes[1].set_title('Validation accuracy per epoch'); axes[1].set_xlabel('Epoch'); axes[1].set_ylabel('Accuracy (%)')
axes[1].set_ylim(93, 99.5); axes[1].yaxis.set_major_formatter(lambda x, _: f'{x:.0f}%')
for ax in axes:
    ax.set_xticks([1, 2, 3]); ax.grid(alpha=0.3); ax.legend()
fig.suptitle('Learning curves — CNN dominates from epoch 1', y=1.03, fontsize=12, fontweight='bold')
…

import seaborn as sns
from sklearn.metrics import confusion_matrix

fig, axes = plt.subplots(1, 2, figsize=(13, 5.2))
for ax, arch in zip(axes, ['mlp', 'cnn']):
    y_true, y_pred = preds[arch]['y_true'], preds[arch]['y_pred']
    cm = confusion_matrix(y_true, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', cbar=False, square=True,
                xticklabels=range(10), yticklabels=range(10), ax=ax, annot_kws={'size': 8})
    acc = (y_pred == y_true).mean()
    ax.set_title(f'{arch.upper()}  —  test acc {acc:.2%}  ({(y_pred != y_true).sum()} errors)', fontsize=11)
    ax.set_xlabel('Predicted'); ax.set_ylabel('True')
fig.suptitle('Confusion matrices on the 10k test set', y=1.02, fontsize=12, fontweight='bold')
fig.tight_layout()
plt.show()

# Per-class accuracy, both architectures
per_class = pd.DataFrame({
…

# Sample predictions from the WINNER: confident correct (top) vs highest-confidence mistakes (bottom)
w = preds[winner]
y_true, y_pred = w['y_true'], w['y_pred']
conf = w['y_prob'].max(axis=1)
imgs = (X_test * STD + MEAN).squeeze(1).numpy()  # un-normalize for display

correct_idx = np.where(y_pred == y_true)[0]
wrong_idx = np.where(y_pred != y_true)[0]
show_correct = correct_idx[np.argsort(conf[correct_idx])[::-1]][:8]          # most confident correct
show_wrong = wrong_idx[np.argsort(conf[wrong_idx])[::-1]][:8]                # most confident mistakes

fig, axes = plt.subplots(2, 8, figsize=(14, 4.2))
for col, idx in enumerate(show_correct):
    axes[0, col].imshow(imgs[idx], cmap='gray')
    axes[0, col].set_title(f'{y_pred[idx]}  ({conf[idx]:.0%})', color='green', fontsize=10)
for col, idx in enumerate(show_wrong):
    axes[1, col].imshow(imgs[idx], cmap='gray')
    axes[1, col].set_title(f'{y_pred[idx]}≠{y_true[idx]} ({conf[idx]:.0%})', color='crimson', fontsize=10)
…

CNN made 129 errors on 10,000 test digits. Mean confidence on errors: 71.5% — the remaining mistakes are genuinely ambiguous handwriting.