train_model.py

from __future__ import annotations

import gzip
import json
import random
import struct
from pathlib import Path
from urllib.request import urlretrieve

import numpy as np
import torch
from PIL import Image, ImageFilter
from torch import nn
from torch.utils.data import DataLoader, Dataset

from mnist_model import DATA_DIR, MNIST_MEAN, MNIST_STD, MODEL_PATH, DigitCNN, center_digit, shift_pixels


METRICS_PATH = Path("model/metrics.json")
MNIST_BASE_URL = "https://storage.googleapis.com/cvdf-datasets/mnist/"
MNIST_FILES = {
    "train_images": "train-images-idx3-ubyte.gz",
    "train_labels": "train-labels-idx1-ubyte.gz",
    "test_images": "t10k-images-idx3-ubyte.gz",
    "test_labels": "t10k-labels-idx1-ubyte.gz",
}


def download_real_mnist() -> None:
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    for filename in MNIST_FILES.values():
        target = DATA_DIR / filename
        if target.exists():
            continue
        print(f"downloading {filename}")
        urlretrieve(MNIST_BASE_URL + filename, target)


def read_images(path: Path) -> np.ndarray:
    with gzip.open(path, "rb") as f:
        magic, count, rows, cols = struct.unpack(">IIII", f.read(16))
        if magic != 2051 or rows != 28 or cols != 28:
            raise ValueError(f"bad image file: {path}")
        data = np.frombuffer(f.read(), dtype=np.uint8)
    return data.reshape(count, 28, 28)


def read_labels(path: Path) -> np.ndarray:
    with gzip.open(path, "rb") as f:
        magic, count = struct.unpack(">II", f.read(8))
        if magic != 2049:
            raise ValueError(f"bad label file: {path}")
        data = np.frombuffer(f.read(), dtype=np.uint8)
    return data.astype(np.int64)


def load_real_mnist() -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray]:
    download_real_mnist()
    x_train = read_images(DATA_DIR / MNIST_FILES["train_images"])
    y_train = read_labels(DATA_DIR / MNIST_FILES["train_labels"])
    x_test = read_images(DATA_DIR / MNIST_FILES["test_images"])
    y_test = read_labels(DATA_DIR / MNIST_FILES["test_labels"])
    return x_train, y_train, x_test, y_test


def augment_canvas_style(image: Image.Image) -> Image.Image:
    angle = random.uniform(-14, 14)
    image = image.rotate(angle, resample=Image.Resampling.BILINEAR, fillcolor=0)

    target_size = random.choice([18, 19, 20, 21, 22, 23, 24])
    image = center_digit(image, target_size=target_size)

    if random.random() < 0.25:
        image = image.filter(ImageFilter.MaxFilter(3))
    if random.random() < 0.15:
        image = image.filter(ImageFilter.GaussianBlur(radius=0.45))

    dy = random.randint(-2, 2)
    dx = random.randint(-2, 2)
    pixels = shift_pixels(np.array(image, dtype=np.uint8), dy, dx)
    return Image.fromarray(pixels, mode="L")


def image_to_tensor(image: Image.Image) -> torch.Tensor:
    pixels = np.array(image, dtype=np.float32) / 255.0
    tensor = torch.from_numpy(pixels).unsqueeze(0)
    return (tensor - MNIST_MEAN) / MNIST_STD


class RealMNISTDataset(Dataset):
    def __init__(self, images: np.ndarray, labels: np.ndarray, augment: bool) -> None:
        self.images = images
        self.labels = labels
        self.augment = augment

    def __len__(self) -> int:
        return len(self.labels)

    def __getitem__(self, index: int) -> tuple[torch.Tensor, torch.Tensor]:
        image = Image.fromarray(self.images[index], mode="L")
        if self.augment:
            image = augment_canvas_style(image)
        else:
            image = center_digit(image, target_size=20)
        return image_to_tensor(image), torch.tensor(self.labels[index], dtype=torch.long)


def evaluate(model: DigitCNN, loader: DataLoader, device: torch.device) -> float:
    model.eval()
    correct = 0
    total = 0
    with torch.no_grad():
        for images, labels in loader:
            images = images.to(device)
            labels = labels.to(device)
            predictions = model(images).argmax(dim=1)
            correct += int((predictions == labels).sum().item())
            total += int(labels.numel())
    return correct / total


def train_model(epochs: int = 5, batch_size: int = 128, learning_rate: float = 1e-3) -> None:
    random.seed(42)
    np.random.seed(42)
    torch.manual_seed(42)

    x_train, y_train, x_test, y_test = load_real_mnist()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"device={device}")
    print(f"train_images={len(y_train)} test_images={len(y_test)}")

    train_loader = DataLoader(
        RealMNISTDataset(x_train, y_train, augment=True),
        batch_size=batch_size,
        shuffle=True,
        num_workers=0,
    )
    test_loader = DataLoader(
        RealMNISTDataset(x_test, y_test, augment=False),
        batch_size=batch_size,
        shuffle=False,
        num_workers=0,
    )

    model = DigitCNN().to(device)
    optimizer = torch.optim.AdamW(model.parameters(), lr=learning_rate, weight_decay=1e-4)
    criterion = nn.CrossEntropyLoss()

    history = []
    best_accuracy = 0.0
    MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)

    for epoch in range(1, epochs + 1):
        model.train()
        running_loss = 0.0
        seen = 0
        for images, labels in train_loader:
            images = images.to(device)
            labels = labels.to(device)
            optimizer.zero_grad()
            loss = criterion(model(images), labels)
            loss.backward()
            optimizer.step()
            running_loss += float(loss.item()) * int(labels.numel())
            seen += int(labels.numel())

        test_accuracy = evaluate(model, test_loader, device)
        train_loss = running_loss / seen
        history.append({"epoch": epoch, "train_loss": train_loss, "test_accuracy": test_accuracy})
        print(f"epoch={epoch} train_loss={train_loss:.4f} test_accuracy={test_accuracy:.4f}")

        if test_accuracy > best_accuracy:
            best_accuracy = test_accuracy
            torch.save(model.cpu().state_dict(), MODEL_PATH)
            model.to(device)

    metrics = {
        "model_type": "PyTorch DigitCNN",
        "training_data": "Real MNIST IDX files downloaded from cvdf-datasets",
        "optimization_strategy": [
            "deeper convolutional neural network",
            "canvas-style augmentation: rotation, scaling, translation, thick strokes, light blur",
            "prediction-time test-time augmentation over multiple target sizes and shifts",
        ],
        "epochs": epochs,
        "batch_size": batch_size,
        "best_test_accuracy": best_accuracy,
        "history": history,
    }
    with open(METRICS_PATH, "w", encoding="utf-8") as f:
        json.dump(metrics, f, indent=2)

    print(f"saved model: {MODEL_PATH}")
    print(f"saved metrics: {METRICS_PATH}")
    print(f"best_test_accuracy={best_accuracy:.4f}")


if __name__ == "__main__":
    train_model()