source/train.py

from __future__ import annotations

import json
import random
import shutil
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import DynamicRoutingCapsuleNet, MatchedMLP, parameter_count
from safetensors.torch import save_file
from torch.nn import functional as F
from torch.utils.data import DataLoader, TensorDataset

PROJECT_DIR = Path(__file__).resolve().parent
ROOT_DIR = PROJECT_DIR.parents[1]
VISION_DATA = ROOT_DIR / "projects" / "tiny-vision-foundry" / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts" / "capsule-pocket"
DATA_DIR = PROJECT_DIR / "data"
SEED = 2179


def load_split(name: str, shuffle: bool) -> DataLoader:
    frame = pd.read_parquet(VISION_DATA / f"{name}.parquet")
    pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 16
    labels = frame["label"].to_numpy(dtype=np.int64, copy=True)
    return DataLoader(
        TensorDataset(torch.from_numpy(pixels), torch.from_numpy(labels)),
        batch_size=128,
        shuffle=shuffle,
        generator=torch.Generator().manual_seed(SEED),
    )


def margin_loss(lengths: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:
    targets = F.one_hot(labels, 10).float()
    positive = targets * F.relu(0.9 - lengths).square()
    negative = 0.5 * (1 - targets) * F.relu(lengths - 0.1).square()
    return (positive + negative).sum(dim=1).mean()


def translate(pixels: torch.Tensor, vertical: int, horizontal: int) -> torch.Tensor:
    images = pixels.reshape(-1, 8, 8)
    shifted = torch.roll(images, shifts=(vertical, horizontal), dims=(1, 2))
    if vertical > 0:
        shifted[:, :vertical] = 0
    elif vertical < 0:
        shifted[:, vertical:] = 0
    if horizontal > 0:
        shifted[:, :, :horizontal] = 0
    elif horizontal < 0:
        shifted[:, :, horizontal:] = 0
    return shifted.reshape(-1, 64)


@torch.inference_mode()
def evaluate(
    model: torch.nn.Module,
    loader: DataLoader,
    *,
    capsule: bool,
    corruption: str,
) -> dict:
    model.eval()
    correct = 0
    total = 0
    for pixels, labels in loader:
        if corruption == "translation":
            variants = [
                translate(pixels, 1, 0),
                translate(pixels, -1, 0),
                translate(pixels, 0, 1),
                translate(pixels, 0, -1),
            ]
            pixels = torch.cat(variants)
            labels = labels.repeat(4)
        elif corruption == "occlusion":
            images = pixels.reshape(-1, 8, 8).clone()
            images[:, 3:5, 3:5] = 0
            pixels = images.reshape(-1, 64)
        scores = model(pixels)[1] if capsule else model(pixels)
        correct += int((scores.argmax(1) == labels).sum())
        total += len(labels)
    return {"accuracy": correct / total, "examples": total}


def train_variant(
    model: torch.nn.Module,
    train_loader: DataLoader,
    validation_loader: DataLoader,
    *,
    capsule: bool,
) -> tuple[dict[str, torch.Tensor], int]:
    optimizer = torch.optim.AdamW(model.parameters(), lr=2e-3, weight_decay=1e-4)
    best = -1.0
    best_epoch = 0
    best_state = None
    for epoch in range(1, 121):
        model.train()
        for pixels, labels in train_loader:
            if capsule:
                _, lengths = model(pixels)
                loss = margin_loss(lengths, labels)
            else:
                loss = F.cross_entropy(model(pixels), labels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
        validation = evaluate(
            model,
            validation_loader,
            capsule=capsule,
            corruption="clean",
        )
        if validation["accuracy"] > best:
            best = validation["accuracy"]
            best_epoch = epoch
            best_state = {
                name: value.detach().cpu().clone()
                for name, value in model.state_dict().items()
            }
        if epoch == 1 or epoch % 10 == 0:
            trackio.log(
                {
                    "variant": "capsule" if capsule else "mlp",
                    "epoch": epoch,
                    "validation_accuracy": validation["accuracy"],
                }
            )
    assert best_state is not None
    return best_state, best_epoch


def main() -> None:
    random.seed(SEED)
    np.random.seed(SEED)
    torch.manual_seed(SEED)
    torch.set_num_threads(1)
    train_loader = load_split("train", True)
    validation_loader = load_split("validation", False)
    test_loader = load_split("test", False)
    capsule = DynamicRoutingCapsuleNet()
    mlp = MatchedMLP()
    assert parameter_count(capsule) == parameter_count(mlp) == 4_060
    trackio.init(
        project="capsule-pocket",
        name="dynamic-routing-digits-v1",
        config={
            "parameters_per_model": 4_060,
            "routing_iterations": capsule.routing_iterations,
            "training_epochs": 120,
        },
    )
    capsule_state, capsule_epoch = train_variant(
        capsule, train_loader, validation_loader, capsule=True
    )
    mlp_state, mlp_epoch = train_variant(
        mlp, train_loader, validation_loader, capsule=False
    )
    capsule.load_state_dict(capsule_state)
    mlp.load_state_dict(mlp_state)
    results = {}
    for name, model, is_capsule, epoch in [
        ("dynamic_routing_capsule", capsule, True, capsule_epoch),
        ("matched_mlp", mlp, False, mlp_epoch),
    ]:
        results[name] = {
            "parameters": parameter_count(model),
            "best_epoch": epoch,
            "clean": evaluate(model, test_loader, capsule=is_capsule, corruption="clean"),
            "one_pixel_translation": evaluate(
                model, test_loader, capsule=is_capsule, corruption="translation"
            ),
            "center_occlusion": evaluate(
                model, test_loader, capsule=is_capsule, corruption="occlusion"
            ),
        }
    report = {
        "experiment": "Dynamic-routing capsule network versus matched MLP",
        "results": results,
    }
    ARTIFACT_DIR.mkdir(parents=True, exist_ok=True)
    DATA_DIR.mkdir(parents=True, exist_ok=True)
    save_file(capsule.state_dict(), ARTIFACT_DIR / "capsule.safetensors")
    save_file(mlp.state_dict(), ARTIFACT_DIR / "matched_mlp.safetensors")
    (ARTIFACT_DIR / "evaluation.json").write_text(
        json.dumps(report, indent=2), encoding="utf-8"
    )
    shutil.copy2(VISION_DATA / "test.parquet", DATA_DIR / "test.parquet")
    trackio.log(
        {
            "capsule_clean_accuracy": results["dynamic_routing_capsule"]["clean"][
                "accuracy"
            ],
            "capsule_translation_accuracy": results["dynamic_routing_capsule"][
                "one_pixel_translation"
            ]["accuracy"],
            "mlp_clean_accuracy": results["matched_mlp"]["clean"]["accuracy"],
            "mlp_translation_accuracy": results["matched_mlp"][
                "one_pixel_translation"
            ]["accuracy"],
        }
    )
    trackio.finish()
    print(json.dumps(report, indent=2))


if __name__ == "__main__":
    main()