Create train_prototype_cell_cifar_10.py

train_prototype_cell_cifar_10.py

@@ -0,0 +1,222 @@
+"""
+SpectralCell Diamond — Maximum Performance Sweep
+===================================================
+Data augmentation + conv spatial head + CE only.
+Sweep D=4 and D=16 back-to-back.
+
+32×32 CIFAR-10 → 16 patches (ps=8) → token_dim=192
+Augmentation: RandomHorizontalFlip + RandomCrop(32, pad=4)
+"""
+
+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')
+
+def extract_patches(images, ps):
+    B, C, H, W = images.shape
+    gh, gw = H // ps, W // ps
+    p = images.reshape(B, C, gh, ps, gw, ps)
+    p = p.permute(0, 2, 4, 1, 3, 5)
+    return p.reshape(B, gh * gw, C * ps * ps), gh, gw
+
+PS = 8
+IMG_SIZE = 32
+GH, GW = IMG_SIZE // PS, IMG_SIZE // PS  # 4, 4
+TOKEN_DIM = 3 * PS * PS  # 192
+N_CLASSES = 10
+CLASSES = ['airplane','auto','bird','cat','deer','dog','frog','horse','ship','truck']
+
+
+class ConvSpatialHead(nn.Module):
+    """Conv on 4×4 grid of formatted tokens.
+    token_dim channels → compress → spatial reasoning → classify.
+    """
+    def __init__(self, token_dim, gh, gw, n_classes=10):
+        super().__init__()
+        self.gh = gh
+        self.gw = gw
+        self.conv = nn.Sequential(
+            nn.Conv2d(token_dim, 64, 3, padding=1),
+            nn.BatchNorm2d(64),
+            nn.GELU(),
+            nn.Conv2d(64, 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, tokens):
+        B = tokens.shape[0]
+        spatial = tokens.permute(0, 2, 1).reshape(B, -1, self.gh, self.gw)
+        h = self.conv(spatial).squeeze(-1).squeeze(-1)
+        return self.head(h)
+
+
+# ── Data with augmentation ───────────────────────────────────────
+
+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=256, shuffle=True, num_workers=4, pin_memory=True, drop_last=True)
+test_loader = torch.utils.data.DataLoader(
+    cifar_test, batch_size=256, shuffle=False, num_workers=4, pin_memory=True)
+
+
+# ── Sweep ────────────────────────────────────────────────────────
+
+CONFIGS = [
+    ("D=4,  V=16, h=128", 16, 4, 128, 2, 1, 2),
+    ("D=16, V=16, h=128", 16, 16, 128, 2, 1, 2),
+]
+
+EPOCHS = 50
+results = {}
+
+for name, V, D, HIDDEN, DEPTH, N_CROSS, N_HEADS in CONFIGS:
+
+    # Adjust n_heads for D
+    if D < N_HEADS * 2:
+        N_HEADS = max(1, D // 2)
+
+    print(f"\n{'═' * 60}")
+    print(f"  {name}")
+    print(f"{'═' * 60}")
+
+    cell = SpectralCell(
+        token_dim=TOKEN_DIM, V=V, D=D,
+        hidden=HIDDEN, depth=DEPTH, n_cross=N_CROSS,
+        n_heads=N_HEADS, max_alpha=0.2,
+    ).to(device)
+
+    clf_head = ConvSpatialHead(TOKEN_DIM, GH, GW, N_CLASSES).to(device)
+
+    cell.summary()
+    n_cell = sum(p.numel() for p in cell.parameters())
+    n_head = sum(p.numel() for p in clf_head.parameters())
+    n_total = n_cell + n_head
+    print(f"  Cell: {n_cell:,}  Head: {n_head:,}  Total: {n_total:,}")
+    print(f"  Augmentation: RandomCrop(32, pad=4) + HFlip")
+
+    all_params = list(cell.parameters()) + list(clf_head.parameters())
+    opt = torch.optim.Adam(all_params, lr=3e-4)
+    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=EPOCHS)
+
+    best_acc = 0
+    t0 = time.time()
+
+    for epoch in range(1, EPOCHS + 1):
+        cell.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)
+            patches, _, _ = extract_patches(images, PS)
+            result = cell.format(patches)
+            logits = clf_head(result['output'])
+            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
+
+        cell.eval()
+        clf_head.eval()
+        val_correct, val_total = 0, 0
+        pcc = torch.zeros(10)
+        pct = torch.zeros(10)
+        last_S_orig, last_S = None, None
+
+        with torch.no_grad():
+            for images, labels in test_loader:
+                images, labels = images.to(device), labels.to(device)
+                patches, _, _ = extract_patches(images, PS)
+                result = cell.format(patches)
+                logits = clf_head(result['output'])
+                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()
+                last_S_orig = result['S_orig']
+                last_S = result['S']
+
+        val_acc = val_correct / val_total
+        star = ''
+        if val_acc > best_acc:
+            best_acc = val_acc
+            star = ' ★'
+
+        S_mean = last_S_orig.mean(dim=(0, 1))
+        erank = cell.effective_rank(last_S_orig.reshape(-1, D)).mean().item()
+        shift = cell.spectral_shift(last_S_orig, last_S)
+
+        if epoch <= 3 or epoch % 5 == 0 or epoch == EPOCHS:
+            s_str = ', '.join(f'{v:.3f}' for v in S_mean.tolist()[:4])
+            if D > 4:
+                s_str += f', ... {S_mean[-1]:.3f}'
+            print(f"    ep{epoch:3d} acc={val_acc:.1%}{star} train={train_acc:.1%} "
+                  f"S=[{s_str}] erank={erank:.2f}")
+
+        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
+    results[name] = (best_acc, n_total, elapsed)
+    print(f"\n  → Best: {best_acc:.1%} | {n_total:,} params | {elapsed:.0f}s")
+
+
+# ── Scoreboard ───────────────────────────────────────────────────
+
+print(f"\n{'═' * 60}")
+print(f"  DIAMOND SCOREBOARD")
+print(f"{'═' * 60}")
+
+print(f"\n  {'Config':<45s} {'Params':>8s} {'Acc':>8s} {'Time':>6s}")
+print(f"  {'-' * 69}")
+
+# Previous results
+prev = [
+    ("Recon only (MSE, no head)", "199K", "—"),
+    ("CE only + mean pool D=4", "212K", "55.1%"),
+    ("MSE+CE + mean pool D=4", "212K", "55.7%"),
+    ("CE only + mean pool D=16", "263K", "56.3%"),
+]
+for n, p, a in prev:
+    print(f"  {n:<45s} {p:>8s} {a:>8s}")
+
+print(f"  {'-' * 69}")
+
+for name, (acc, params, elapsed) in sorted(results.items(), key=lambda x: x[1][0]):
+    print(f"  {name + ' + conv + aug':<45s} {params:>8,} {acc:>7.1%} {elapsed:>5.0f}s")
+
+best = max(results.items(), key=lambda x: x[1][0])
+print(f"\n  Best: {best[0]} → {best[1][0]:.1%}")