MNIST Convolutional Neural Network with PyTorch#
This complete workflow belongs to Convolutional Neural Networks. Keras and PyTorch use the same seeded indices, normalization, architecture intent, training budget, metrics, and diagnostic figures.
In Colab, select Runtime → Run all; the canonical dataset is downloaded by
the framework. To run the example more quickly, set EPOCHS to 1 or 2 in the
data-loading cell.
Runtime dependency check#
import importlib.util
import subprocess
import sys
from pathlib import Path
REQUIRED_RUNTIME = {'torch': 'torch'}
COLAB_EXTRAS = {}
missing_required = [
package for module, package in REQUIRED_RUNTIME.items()
if importlib.util.find_spec(module) is None
]
missing_extras = [
package for module, package in COLAB_EXTRAS.items()
if importlib.util.find_spec(module) is None
]
if missing_extras and "google.colab" in sys.modules:
subprocess.check_call(
[sys.executable, "-m", "pip", "install", "-q", *missing_extras]
)
missing_extras = []
if missing_required or missing_extras:
missing = ", ".join(missing_required + missing_extras)
raise RuntimeError(
f"Missing notebook dependencies: {missing}. Locally run "
"`uv sync --group notebooks`; in Colab restart the runtime if an "
"installation cell just changed the environment."
)
print("runtime dependency check passed")
runtime dependency check passed
%matplotlib inline
Imports and reproducibility#
import hashlib
import json
import os
import random
from pathlib import Path
import matplotlib.pyplot as plt
import numpy as np
import torch
from sklearn.metrics import (
ConfusionMatrixDisplay,
accuracy_score,
classification_report,
confusion_matrix,
)
from sklearn.model_selection import train_test_split
from torch import nn
from torch.utils.data import DataLoader, TensorDataset
from torchvision import datasets
print(f"PyTorch {torch.__version__}")
PyTorch 2.13.0
Inspect class coverage#
training_counts = np.bincount(y_train, minlength=10)
fig, ax = plt.subplots(figsize=(8, 3.2))
ax.bar(np.arange(10), training_counts)
ax.set(
title="Training-set class distribution",
xlabel="Class",
ylabel="Samples",
xticks=np.arange(10),
)
plt.show()
Define the model#
class ImageClassifier(nn.Module):
def __init__(self):
super().__init__()
self.features = nn.Sequential(
nn.Conv2d(1, 32, 3), nn.ReLU(), nn.MaxPool2d(2),
nn.Conv2d(32, 64, 3), nn.ReLU(), nn.MaxPool2d(2), nn.Dropout(0.25),
)
self.classifier = nn.Sequential(
nn.Flatten(), nn.Linear(64 * 5 * 5, 200), nn.ReLU(),
nn.Linear(200, 150), nn.ReLU(), nn.Dropout(0.5), nn.Linear(150, 10),
)
def forward(self, values):
return self.classifier(self.features(values))
model = ImageClassifier()
print(model)
ImageClassifier(
(features): Sequential(
(0): Conv2d(1, 32, kernel_size=(3, 3), stride=(1, 1))
(1): ReLU()
(2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(3): Conv2d(32, 64, kernel_size=(3, 3), stride=(1, 1))
(4): ReLU()
(5): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
(6): Dropout(p=0.25, inplace=False)
)
(classifier): Sequential(
(0): Flatten(start_dim=1, end_dim=-1)
(1): Linear(in_features=1600, out_features=200, bias=True)
(2): ReLU()
(3): Linear(in_features=200, out_features=150, bias=True)
(4): ReLU()
(5): Dropout(p=0.5, inplace=False)
(6): Linear(in_features=150, out_features=10, bias=True)
)
)
Train with validation evidence#
def tensor_dataset(images, labels):
return TensorDataset(torch.from_numpy(images), torch.from_numpy(labels).long())
generator = torch.Generator().manual_seed(SEED)
train_loader = DataLoader(
tensor_dataset(x_train, y_train),
batch_size=BATCH_SIZE,
shuffle=True,
generator=generator,
)
validation_loader = DataLoader(
tensor_dataset(x_validation, y_validation), batch_size=BATCH_SIZE
)
test_loader = DataLoader(tensor_dataset(x_test, y_test), batch_size=BATCH_SIZE)
model = model.to(DEVICE)
optimizer = torch.optim.Adam(model.parameters())
loss_function = nn.CrossEntropyLoss()
history = {"loss": [], "val_loss": [], "accuracy": [], "val_accuracy": []}
def epoch_pass(loader, training):
model.train(training)
total_loss = 0.0
correct = 0
for features, target in loader:
features, target = features.to(DEVICE), target.to(DEVICE)
if training:
optimizer.zero_grad()
logits = model(features)
loss = loss_function(logits, target)
if training:
loss.backward()
optimizer.step()
total_loss += loss.item() * len(target)
correct += (logits.argmax(1) == target).sum().item()
return total_loss / len(loader.dataset), correct / len(loader.dataset)
for epoch in range(EPOCHS):
train_loss, train_accuracy = epoch_pass(train_loader, True)
with torch.no_grad():
validation_loss, validation_accuracy = epoch_pass(validation_loader, False)
history["loss"].append(train_loss)
history["accuracy"].append(train_accuracy)
history["val_loss"].append(validation_loss)
history["val_accuracy"].append(validation_accuracy)
print(
f"epoch {epoch + 1}: loss={train_loss:.4f}, accuracy={train_accuracy:.4f}, "
f"val_loss={validation_loss:.4f}, val_accuracy={validation_accuracy:.4f}"
)
fig, axes = plt.subplots(1, 2, figsize=(10, 3.5))
axes[0].plot(history["loss"], marker="o", label="training")
axes[0].plot(history["val_loss"], marker="o", label="validation")
axes[0].set(title="Cross-entropy loss", xlabel="Epoch")
axes[1].plot(history["accuracy"], marker="o", label="training")
axes[1].plot(history["val_accuracy"], marker="o", label="validation")
axes[1].set(title="Accuracy", xlabel="Epoch")
for axis in axes:
axis.legend()
axis.grid(alpha=0.25)
plt.tight_layout()
plt.show()
epoch 1: loss=0.5572, accuracy=0.8216, val_loss=0.1168, val_accuracy=0.9658
epoch 2: loss=0.1252, accuracy=0.9631, val_loss=0.0731, val_accuracy=0.9775
epoch 3: loss=0.0830, accuracy=0.9760, val_loss=0.0564, val_accuracy=0.9837
epoch 4: loss=0.0659, accuracy=0.9803, val_loss=0.0546, val_accuracy=0.9843
epoch 5: loss=0.0528, accuracy=0.9846, val_loss=0.0430, val_accuracy=0.9880
Evaluate once on the test set#
model.eval()
prediction_parts = []
with torch.no_grad():
for features, _ in test_loader:
prediction_parts.append(model(features.to(DEVICE)).argmax(1).cpu().numpy())
test_predictions = np.concatenate(prediction_parts)
test_accuracy = accuracy_score(y_test, test_predictions)
cm = confusion_matrix(y_test, test_predictions, labels=np.arange(10))
print(f"test accuracy: {test_accuracy:.4f}")
print(
classification_report(
y_test,
test_predictions,
labels=np.arange(10),
digits=3,
zero_division=0,
)
)
print(
"HELIO_RESULT "
+ json.dumps(
{
"split_signature": split_signature,
"test_accuracy": float(test_accuracy),
"confusion_shape": list(cm.shape),
},
sort_keys=True,
)
)
assert cm.shape == (10, 10)
fig, ax = plt.subplots(figsize=(7, 6))
ConfusionMatrixDisplay(cm, display_labels=np.arange(10)).plot(
ax=ax, colorbar=False, values_format="d"
)
ax.set_title("Test confusion matrix")
plt.show()
mistakes = np.flatnonzero(test_predictions != y_test)[:12]
if len(mistakes):
fig, axes = plt.subplots(3, 4, figsize=(9, 7))
for axis, index in zip(axes.flat, mistakes):
image = x_test[index]
if image.shape[0] in (1, 3):
image = np.transpose(image, (1, 2, 0))
axis.imshow(image.squeeze(), cmap="gray" if image.squeeze().ndim == 2 else None)
axis.set_title(f"true={y_test[index]}, pred={test_predictions[index]}")
axis.axis("off")
plt.tight_layout()
plt.show()
test accuracy: 0.9900
precision recall f1-score support
0 0.986 0.995 0.990 980
1 0.991 1.000 0.996 1135
2 0.991 0.995 0.993 1032
3 0.990 0.991 0.991 1010
4 0.989 0.996 0.992 982
5 0.989 0.988 0.988 892
6 0.995 0.991 0.993 958
7 0.985 0.988 0.987 1028
8 0.988 0.985 0.986 974
9 0.996 0.970 0.983 1009
accuracy 0.990 10000
macro avg 0.990 0.990 0.990 10000
weighted avg 0.990 0.990 0.990 10000
HELIO_RESULT {"confusion_shape": [10, 10], "split_signature": "829d1431018b38bc", "test_accuracy": 0.99}
Try it yourself#
Change one choice at a time and keep the data split and evaluation unchanged:
change the first convolution from 32 to 64 filters.
compare 3×3 and 5×5 kernels.
remove one pooling operation and inspect the tensor shapes.