source/train.py

from __future__ import annotations

import json
import random
from pathlib import Path

import numpy as np
import pandas as pd
import torch
import trackio
from model import (
    DenseControl,
    PocketMoE,
    active_parameter_count,
    parameter_count,
)
from safetensors.torch import save_file
from sklearn.metrics import accuracy_score, f1_score
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]
DATA_DIR = ROOT_DIR / "projects" / "tiny-vision-foundry" / "data"
ARTIFACT_DIR = PROJECT_DIR / "artifacts"


def seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)


def load_split(name: str) -> tuple[torch.Tensor, torch.Tensor]:
    frame = pd.read_parquet(DATA_DIR / f"{name}.parquet")
    pixels = np.stack(frame["image"].to_numpy()).astype(np.float32) / 16.0
    labels = frame["label"].to_numpy(dtype=np.int64, copy=True)
    return torch.from_numpy(pixels), torch.from_numpy(labels)


@torch.inference_mode()
def evaluate_dense(
    model: DenseControl,
    pixels: torch.Tensor,
    labels: torch.Tensor,
) -> dict:
    model.eval()
    predictions = model(pixels).argmax(dim=1).numpy()
    return {
        "accuracy": float(accuracy_score(labels.numpy(), predictions)),
        "macro_f1": float(f1_score(labels.numpy(), predictions, average="macro")),
    }


@torch.inference_mode()
def evaluate_moe(
    model: PocketMoE,
    pixels: torch.Tensor,
    labels: torch.Tensor,
) -> dict:
    model.eval()
    logits, router_probabilities, sparse_weights = model(pixels)
    predictions = logits.argmax(dim=1).numpy()
    utilization = sparse_weights.mean(dim=0).numpy()
    dominant_by_class = {}
    for label in range(10):
        selected = labels == label
        class_utilization = sparse_weights[selected].mean(dim=0)
        dominant_by_class[str(label)] = {
            "expert": int(class_utilization.argmax()),
            "routing_share": float(class_utilization.max()),
        }
    entropy = -(
        router_probabilities * torch.log(torch.clamp(router_probabilities, min=1e-9))
    ).sum(dim=1)
    return {
        "accuracy": float(accuracy_score(labels.numpy(), predictions)),
        "macro_f1": float(f1_score(labels.numpy(), predictions, average="macro")),
        "expert_utilization": utilization.tolist(),
        "utilization_coefficient_of_variation": float(
            utilization.std() / utilization.mean()
        ),
        "mean_router_entropy": float(entropy.mean()),
        "maximum_router_entropy": float(np.log(model.expert_count)),
        "dominant_expert_by_digit": dominant_by_class,
    }


def train_dense(
    train_pixels: torch.Tensor,
    train_labels: torch.Tensor,
    validation_pixels: torch.Tensor,
    validation_labels: torch.Tensor,
) -> tuple[DenseControl, dict]:
    seed_everything(2042)
    model = DenseControl()
    loader = DataLoader(
        TensorDataset(train_pixels, train_labels),
        batch_size=64,
        shuffle=True,
        generator=torch.Generator().manual_seed(2042),
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0025, weight_decay=0.002)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=90)
    best_accuracy = -1.0
    best_epoch = 0
    best_state = None
    for epoch in range(1, 91):
        model.train()
        for pixels, labels in loader:
            loss = F.cross_entropy(model(pixels), labels)
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
        scheduler.step()
        validation = evaluate_dense(model, validation_pixels, validation_labels)
        if validation["accuracy"] > best_accuracy:
            best_accuracy = validation["accuracy"]
            best_epoch = epoch
            best_state = {
                key: value.detach().cpu().clone()
                for key, value in model.state_dict().items()
            }
    assert best_state is not None
    model.load_state_dict(best_state)
    return model, {
        "best_epoch": best_epoch,
        "best_validation_accuracy": best_accuracy,
    }


def train_moe(
    train_pixels: torch.Tensor,
    train_labels: torch.Tensor,
    validation_pixels: torch.Tensor,
    validation_labels: torch.Tensor,
) -> tuple[PocketMoE, dict]:
    seed_everything(2042)
    model = PocketMoE()
    loader = DataLoader(
        TensorDataset(train_pixels, train_labels),
        batch_size=64,
        shuffle=True,
        generator=torch.Generator().manual_seed(2042),
    )
    optimizer = torch.optim.AdamW(model.parameters(), lr=0.0025, weight_decay=0.002)
    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=90)
    best_accuracy = -1.0
    best_epoch = 0
    best_state = None
    for epoch in range(1, 91):
        model.train()
        losses = []
        for pixels, labels in loader:
            logits, router_probabilities, _ = model(pixels)
            classification = F.cross_entropy(logits, labels)
            importance = router_probabilities.mean(dim=0)
            balance = ((importance * model.expert_count - 1) ** 2).mean()
            loss = classification + 0.025 * balance
            optimizer.zero_grad(set_to_none=True)
            loss.backward()
            optimizer.step()
            losses.append(loss.item())
        scheduler.step()
        validation = evaluate_moe(model, validation_pixels, validation_labels)
        if validation["accuracy"] > best_accuracy:
            best_accuracy = validation["accuracy"]
            best_epoch = epoch
            best_state = {
                key: value.detach().cpu().clone()
                for key, value in model.state_dict().items()
            }
        if epoch == 1 or epoch % 10 == 0:
            trackio.log(
                {
                    "epoch": epoch,
                    "moe_train_loss": float(np.mean(losses)),
                    "moe_validation_accuracy": validation["accuracy"],
                    "moe_utilization_cv": validation[
                        "utilization_coefficient_of_variation"
                    ],
                    "learning_rate": scheduler.get_last_lr()[0],
                }
            )
    assert best_state is not None
    model.load_state_dict(best_state)
    return model, {
        "best_epoch": best_epoch,
        "best_validation_accuracy": best_accuracy,
    }


def main() -> None:
    train_pixels, train_labels = load_split("train")
    validation_pixels, validation_labels = load_split("validation")
    test_pixels, test_labels = load_split("test")
    trackio.init(
        project="pocket-moe",
        name="top2-versus-dense-v1",
        config={
            "experts": 4,
            "active_experts": 2,
            "moe_parameters": parameter_count(PocketMoE()),
            "dense_parameters": parameter_count(DenseControl()),
        },
    )
    dense, dense_training = train_dense(
        train_pixels,
        train_labels,
        validation_pixels,
        validation_labels,
    )
    moe, moe_training = train_moe(
        train_pixels,
        train_labels,
        validation_pixels,
        validation_labels,
    )
    dense_test = evaluate_dense(dense, test_pixels, test_labels)
    moe_test = evaluate_moe(moe, test_pixels, test_labels)
    results = {
        "model": "Pocket MoE",
        "experts": moe.expert_count,
        "active_experts_per_example": moe.top_k,
        "stored_parameters": parameter_count(moe),
        "active_parameters_per_example": active_parameter_count(moe),
        "dense_control_parameters": parameter_count(dense),
        "moe_training": moe_training,
        "dense_training": dense_training,
        "moe_test": moe_test,
        "dense_control_test": dense_test,
        "accuracy_delta": moe_test["accuracy"] - dense_test["accuracy"],
    }
    trackio.log(
        {
            "moe_test_accuracy": moe_test["accuracy"],
            "dense_test_accuracy": dense_test["accuracy"],
            "moe_test_utilization_cv": moe_test["utilization_coefficient_of_variation"],
        }
    )
    trackio.finish()
    moe_dir = ARTIFACT_DIR / "pocket-moe"
    dense_dir = ARTIFACT_DIR / "dense-control"
    moe_dir.mkdir(parents=True, exist_ok=True)
    dense_dir.mkdir(parents=True, exist_ok=True)
    save_file(moe.state_dict(), moe_dir / "model.safetensors")
    save_file(dense.state_dict(), dense_dir / "model.safetensors")
    (moe_dir / "evaluation.json").write_text(
        json.dumps(results, indent=2),
        encoding="utf-8",
    )
    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()