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