MNIST CNN Tuning with Optuna and PyTorch

This complete workflow treats tuning as a validation experiment. Both frameworks use the same seed-42 split, discrete search space, four-trial budget, ten-epoch maximum, and validation accuracy objective.

To run the search more quickly, reduce TRIALS or SEARCH_EPOCHS below.

Runtime dependency check

import importlib.util
import subprocess
import sys
from pathlib import Path

REQUIRED_RUNTIME = {'torch': 'torch'}
COLAB_EXTRAS = {'optuna': 'optuna'}
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 shared split

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__}")

import importlib.util
import subprocess
import sys
if importlib.util.find_spec("optuna") is None:
    subprocess.check_call(
        [sys.executable, "-m", "pip", "install", "-q", "optuna>=4,<5"]
    )
import warnings

warnings.filterwarnings("ignore", message="IProgress not found.*")
import optuna
optuna.logging.set_verbosity(optuna.logging.WARNING)

PyTorch 2.13.0

SEED = 42
EPOCHS = 5  # Reduce to 1 or 2 for a quicker run.
BATCH_SIZE = 128

random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.use_deterministic_algorithms(True)
DEVICE = torch.device(
    "cuda" if torch.cuda.is_available()
    else "mps" if torch.backends.mps.is_available()
    else "cpu"
)


data_root = Path(
    os.getenv("HELIO_DATA_DIR", Path.home() / ".cache" / "helio-data-methods")
)
development_dataset = datasets.MNIST(data_root, train=True, download=True)
test_dataset = datasets.MNIST(data_root, train=False, download=True)
x_development = development_dataset.data.numpy()
y_development = development_dataset.targets.numpy()
x_test = test_dataset.data.numpy()
y_test = test_dataset.targets.numpy()

all_indices = np.arange(len(y_development))
train_indices, validation_indices = train_test_split(
    all_indices,
    test_size=10000,
    random_state=SEED,
    stratify=y_development,
)
split_signature = hashlib.sha256(
    validation_indices.astype("<i8").tobytes()
).hexdigest()[:16]

x_train = x_development[train_indices].astype("float32") / 255.0
y_train = y_development[train_indices].astype(np.int64)
x_validation = x_development[validation_indices].astype("float32") / 255.0
y_validation = y_development[validation_indices].astype(np.int64)
x_test = x_test.astype("float32") / 255.0
y_test = y_test.astype(np.int64)
x_train = x_train[:, np.newaxis, ...]
x_validation = x_validation[:, np.newaxis, ...]
x_test = x_test[:, np.newaxis, ...]

assert set(train_indices).isdisjoint(validation_indices)
print(f"split signature: {split_signature}")
print(
    f"train={len(y_train):,}, validation={len(y_validation):,}, "
    f"test={len(y_test):,}, epochs={EPOCHS}"
)

split signature: 829d1431018b38bc
train=50,000, validation=10,000, test=10,000, epochs=5

Fixed search contract

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()

TRIALS = 4  # Reduce to 2 for a quicker search.
SEARCH_EPOCHS = 10  # Reduce to 1 or 2 for quicker trials.
SEARCH_SPACE = {
    "filters_1": [16, 32],
    "filters_2": [32, 64],
    "dense_units": [100, 200],
    "learning_rate": [1e-2, 1e-3, 1e-4],
}

Run Optuna

def build_trial_model(trial):
    filters_1 = trial.suggest_categorical("filters_1", SEARCH_SPACE["filters_1"])
    filters_2 = trial.suggest_categorical("filters_2", SEARCH_SPACE["filters_2"])
    dense_units = trial.suggest_categorical("dense_units", SEARCH_SPACE["dense_units"])
    return nn.Sequential(
        nn.Conv2d(1, filters_1, 3), nn.ReLU(), nn.MaxPool2d(2),
        nn.Conv2d(filters_1, filters_2, 3), nn.ReLU(), nn.MaxPool2d(2),
        nn.Dropout(0.25), nn.Flatten(), nn.Linear(filters_2 * 5 * 5, dense_units),
        nn.ReLU(), nn.Dropout(0.5), nn.Linear(dense_units, 10),
    )


train_dataset = TensorDataset(torch.from_numpy(x_train), torch.from_numpy(y_train).long())
validation_features = torch.from_numpy(x_validation)
validation_targets = torch.from_numpy(y_validation).long()


def objective(trial):
    torch.manual_seed(SEED + trial.number)
    candidate = build_trial_model(trial)
    learning_rate = trial.suggest_categorical(
        "learning_rate", SEARCH_SPACE["learning_rate"]
    )
    candidate_optimizer = torch.optim.Adam(candidate.parameters(), lr=learning_rate)
    criterion = nn.CrossEntropyLoss()
    loader = DataLoader(
        train_dataset,
        batch_size=BATCH_SIZE,
        shuffle=True,
        generator=torch.Generator().manual_seed(SEED + trial.number),
    )
    best_accuracy = 0.0
    for epoch in range(SEARCH_EPOCHS):
        candidate.train()
        for features, target in loader:
            candidate_optimizer.zero_grad()
            loss = criterion(candidate(features), target)
            loss.backward()
            candidate_optimizer.step()
        candidate.eval()
        with torch.no_grad():
            accuracy = (
                candidate(validation_features).argmax(1) == validation_targets
            ).float().mean().item()
        best_accuracy = max(best_accuracy, accuracy)
        trial.report(accuracy, epoch)
    return best_accuracy


study = optuna.create_study(direction="maximize", sampler=optuna.samplers.TPESampler(seed=SEED))
study.optimize(objective, n_trials=TRIALS)
print("best parameters:", study.best_params)

best parameters: {'filters_1': 32, 'filters_2': 64, 'dense_units': 100, 'learning_rate': 0.001}

Rebuild the selected model and evaluate once

class FixedTrial:
    def suggest_categorical(self, name, values):
        return study.best_params[name]


model = build_trial_model(FixedTrial())

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(), lr=study.best_params["learning_rate"])
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.4181, accuracy=0.8712, val_loss=0.1042, val_accuracy=0.9694

epoch 2: loss=0.1480, accuracy=0.9568, val_loss=0.0757, val_accuracy=0.9781

epoch 3: loss=0.1119, accuracy=0.9672, val_loss=0.0596, val_accuracy=0.9833

epoch 4: loss=0.0930, accuracy=0.9730, val_loss=0.0517, val_accuracy=0.9862

epoch 5: loss=0.0790, accuracy=0.9765, val_loss=0.0469, val_accuracy=0.9864

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.9886
              precision    recall  f1-score   support

           0      0.987     0.998     0.992       980
           1      0.990     0.998     0.994      1135
           2      0.990     0.984     0.987      1032
           3      0.990     0.988     0.989      1010
           4      0.988     0.998     0.993       982
           5      0.989     0.987     0.988       892
           6      0.998     0.991     0.994       958
           7      0.970     0.991     0.980      1028
           8      0.994     0.978     0.986       974
           9      0.992     0.972     0.982      1009

    accuracy                          0.989     10000
   macro avg      0.989     0.988     0.989     10000
weighted avg      0.989     0.989     0.989     10000

HELIO_RESULT {"confusion_shape": [10, 10], "split_signature": "829d1431018b38bc", "test_accuracy": 0.9886}

Try it yourself

Change one choice at a time and keep the data split and evaluation unchanged:

• add dropout values to the search space.

• increase the number of trials and inspect whether the result is stable.

• repeat the search with a second sampler seed.