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