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