Create train_prototype_conduit_battery_cifar_10.py

train_prototype_conduit_battery_cifar_10.py

@@ -0,0 +1,257 @@
+"""
+Train ConduitBattery — CIFAR-10 Classification
+=================================================
+Conv pathway: (B, 3, 32, 32) → 64 groups of 16×16 local matrices
+SVD with conduit telemetry on each group.
+Two-stream relational processing: geometric + content with FiLM.
+Classify from global_token (pooled geometric summary).
+
+Config: input_dim=3, rank=16, conv_window=4, geom_dim=32, model_dim=64
+CE loss, augmentation.
+"""
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import time
+from tqdm import tqdm
+import torchvision
+import torchvision.transforms as T
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+# ── Config ───────────────────────────────────────────────────────
+
+N_CLASSES = 10
+EPOCHS = 50
+BATCH = 256
+LR = 3e-4
+CLASSES = ['airplane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
+
+battery_config = BaseConfig(
+    input_dim=3,
+    model_dim=64,
+    out_dim=3,
+    rank=16,
+    geom_dim=32,
+    relation_depth=2,
+    num_heads=4,
+    mlp_ratio=2.0,
+    dropout=0.0,
+    conv_window=4,
+    max_spectral_shift=0.20,
+    max_scale_shift=0.10,
+    residual_init=0.0,
+    use_row_norm=True,
+    use_conduit=True,
+    compute_dtype="fp64",
+)
+
+# ── Classification Heads ─────────────────────────────────────────
+
+class GlobalHead(nn.Module):
+    """Classify from global_token (B, model_dim) → logits."""
+    def __init__(self, model_dim, n_classes=10):
+        super().__init__()
+        self.head = nn.Sequential(
+            nn.LayerNorm(model_dim),
+            nn.Linear(model_dim, n_classes),
+        )
+    def forward(self, global_token):
+        return self.head(global_token)
+
+
+class GridConvHead(nn.Module):
+    """Classify from analysis_grid (B, model_dim, gh, gw) via small conv."""
+    def __init__(self, model_dim, n_classes=10):
+        super().__init__()
+        self.conv = nn.Sequential(
+            nn.Conv2d(model_dim, 64, 3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.GELU(),
+            nn.Conv2d(64, 32, 3, padding=1),
+            nn.BatchNorm2d(32),
+            nn.GELU(),
+            nn.AdaptiveAvgPool2d(1),
+        )
+        self.head = nn.Linear(32, n_classes)
+
+    def forward(self, grid):
+        h = self.conv(grid).squeeze(-1).squeeze(-1)
+        return self.head(h)
+
+
+# ── Data ─────────────────────────────────────────────────────────
+
+train_transform = T.Compose([
+    T.RandomCrop(32, padding=4),
+    T.RandomHorizontalFlip(),
+    T.ToTensor(),
+])
+test_transform = T.Compose([T.ToTensor()])
+
+cifar_train = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=train_transform)
+cifar_test = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=test_transform)
+
+train_loader = torch.utils.data.DataLoader(
+    cifar_train, batch_size=BATCH, shuffle=True, num_workers=4, pin_memory=True, drop_last=True)
+test_loader = torch.utils.data.DataLoader(
+    cifar_test, batch_size=BATCH, shuffle=False, num_workers=4, pin_memory=True)
+
+
+# ── Build Model ──────────────────────────────────────────────────
+
+battery = ConduitBattery(battery_config).to(device)
+head_global = GlobalHead(battery_config.model_dim, N_CLASSES).to(device)
+head_grid = GridConvHead(battery_config.model_dim, N_CLASSES).to(device)
+
+n_battery = sum(p.numel() for p in battery.parameters())
+n_global = sum(p.numel() for p in head_global.parameters())
+n_grid = sum(p.numel() for p in head_grid.parameters())
+
+print(f"ConduitBattery config:")
+print(f"  input_dim={battery_config.input_dim}, rank={battery_config.rank}")
+print(f"  conv_window={battery_config.conv_window} → 8×8=64 groups of 16×16")
+print(f"  geom_dim={battery_config.geom_dim}, model_dim={battery_config.model_dim}")
+print(f"  relation_depth={battery_config.relation_depth}")
+print(f"  conduit={'ON' if battery_config.use_conduit else 'OFF'}")
+print(f"\nParams:")
+print(f"  Battery:    {n_battery:,}")
+print(f"  GlobalHead: {n_global:,}")
+print(f"  GridHead:   {n_grid:,}")
+
+
+# ═══════════════════════════════════════════════════════════════
+# EXPERIMENT A: global_token → classification
+# ═══════════════════════════════════════════════════════════════
+
+def run_experiment(battery, clf_head, head_name, get_logits_fn):
+    """Train and evaluate one configuration."""
+    all_params = list(battery.parameters()) + list(clf_head.parameters())
+    n_total = sum(p.numel() for p in all_params)
+    opt = torch.optim.Adam(all_params, lr=LR)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=EPOCHS)
+
+    print(f"\n{'═' * 60}")
+    print(f"  {head_name} — {n_total:,} params")
+    print(f"{'═' * 60}")
+
+    best_acc = 0
+    t0 = time.time()
+
+    for epoch in range(1, EPOCHS + 1):
+        battery.train()
+        clf_head.train()
+        correct, total = 0, 0
+
+        for images, labels in tqdm(train_loader, desc=f"Ep {epoch:2d}", leave=False):
+            images, labels = images.to(device), labels.to(device)
+            result = battery(images, input_kind='conv', return_state=False)
+            logits = get_logits_fn(result, clf_head)
+            loss = F.cross_entropy(logits, labels)
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+            correct += (logits.argmax(-1) == labels).sum().item()
+            total += images.shape[0]
+
+        sched.step()
+        train_acc = correct / total
+
+        battery.eval()
+        clf_head.eval()
+        val_correct, val_total = 0, 0
+        pcc = torch.zeros(10)
+        pct = torch.zeros(10)
+
+        with torch.no_grad():
+            for images, labels in test_loader:
+                images, labels = images.to(device), labels.to(device)
+                result = battery(images, input_kind='conv')
+                logits = get_logits_fn(result, clf_head)
+                preds = logits.argmax(-1)
+                val_correct += (preds == labels).sum().item()
+                val_total += images.shape[0]
+                for c in range(10):
+                    m = labels == c
+                    pcc[c] += (preds[m] == labels[m]).sum().item()
+                    pct[c] += m.sum().item()
+
+        val_acc = val_correct / val_total
+        star = ''
+        if val_acc > best_acc:
+            best_acc = val_acc
+            star = ' ★'
+
+        if epoch <= 3 or epoch % 5 == 0 or epoch == EPOCHS:
+            S = result['S']
+            S_mean = S.mean(dim=(0, 1))
+            s_str = ', '.join(f'{v:.3f}' for v in S_mean.tolist()[:4])
+            if S.shape[-1] > 4:
+                s_str += f', ... {S_mean[-1]:.3f}'
+            erank = result['effective_rank'].mean().item()
+            entropy = result['spectral_entropy'].mean().item()
+            shift = (result['S_shifted'] - result['S']).abs().mean().item()
+
+            print(f"    ep{epoch:3d} acc={val_acc:.1%}{star} train={train_acc:.1%} "
+                  f"S=[{s_str}] erank={erank:.2f} shift={shift:.5f}")
+
+        if epoch <= 2 or epoch % 10 == 0 or epoch == EPOCHS:
+            pca = pcc / (pct + 1e-8)
+            print(f"    Per-class:")
+            for c in range(10):
+                bar = '█' * int(pca[c] * 20)
+                print(f"      {CLASSES[c]:<10s} {pca[c]:5.1%} {bar}")
+
+    elapsed = time.time() - t0
+    print(f"\n  → {head_name}: Best={best_acc:.1%} | {n_total:,} params | {elapsed:.0f}s")
+    return best_acc, n_total, elapsed
+
+
+# Run experiments
+results = {}
+
+# A: global_token → Linear
+battery_a = ConduitBattery(battery_config).to(device)
+head_a = GlobalHead(battery_config.model_dim, N_CLASSES).to(device)
+acc_a, params_a, time_a = run_experiment(
+    battery_a, head_a, "ConduitBattery + GlobalHead",
+    lambda r, h: h(r['global_token'])
+)
+results['global'] = (acc_a, params_a, time_a)
+
+# B: analysis_grid → Conv
+battery_b = ConduitBattery(battery_config).to(device)
+head_b = GridConvHead(battery_config.model_dim, N_CLASSES).to(device)
+acc_b, params_b, time_b = run_experiment(
+    battery_b, head_b, "ConduitBattery + GridConvHead",
+    lambda r, h: h(r['analysis_grid'])
+)
+results['grid_conv'] = (acc_b, params_b, time_b)
+
+
+# ── Scoreboard ───────────────────────────────────────────────────
+
+print(f"\n{'═' * 60}")
+print(f"  SCOREBOARD")
+print(f"{'═' * 60}")
+
+print(f"\n  {'Config':<50s} {'Params':>8s} {'Acc':>8s}")
+print(f"  {'-' * 68}")
+
+# SpectralCell baselines
+prev = [
+    ("SpectralCell CE+mean D=4", "212K", "55.1%"),
+    ("SpectralCell CE+mean D=16", "263K", "56.3%"),
+    ("SpectralCell CE+conv D=4 +aug", "366K", "75.7%"),
+    ("SpectralCell CE+conv D=16 +aug", "416K", "76.0%"),
+]
+for n, p, a in prev:
+    print(f"  {n:<50s} {p:>8s} {a:>8s}")
+
+print(f"  {'-' * 68}")
+for name, (acc, params, elapsed) in results.items():
+    print(f"  {'ConduitBattery + ' + name:<50s} {params:>8,} {acc:>7.1%}")
+
+best = max(results.items(), key=lambda x: x[1][0])
+print(f"\n  Best: ConduitBattery + {best[0]} → {best[1][0]:.1%}")