Create cell3_trainer.py

cell3_trainer.py

@@ -0,0 +1,246 @@
+"""
+Cell 3: Trainer — Patch Cross-Attention Shape Classifier (8×16×16)
+===================================================================
+Run after Cell 1 (generator) and Cell 2 (model).
+Everything from prior cells is already in kernel scope.
+"""
+
+import os, time, math
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import TensorDataset, DataLoader
+
+
+# === Augmentation (vectorized for 8×16×16) ====================================
+
+def deform_grid(grid, p_dropout=0.05, p_add=0.05, p_shift=0.08):
+    """Vectorized voxel augmentation for 8×16×16 grids."""
+    B = grid.shape[0]
+    device = grid.device
+    r = torch.rand(B, 3, device=device)
+    out = grid.clone()
+
+    # Voxel dropout
+    drop_sel = (r[:, 0] < p_dropout).view(B, 1, 1, 1)
+    keep = torch.rand_like(out) > 0.10
+    out = torch.where(drop_sel, out * keep.float(), out)
+
+    # Boundary addition
+    add_sel = (r[:, 1] < p_add).view(B, 1, 1, 1).float()
+    dilated = F.max_pool3d(out.unsqueeze(1), kernel_size=3, stride=1, padding=1).squeeze(1)
+    boundary = ((dilated > 0.5) & (out < 0.5)).float()
+    add_noise = (torch.rand_like(out) < 0.2).float()
+    out = (out + boundary * add_noise * add_sel).clamp(max=1.0)
+
+    # Small translation (1 voxel — grid is small)
+    shift_sel = (r[:, 2] < p_shift)
+    axes = torch.randint(3, (B,), device=device)
+    dirs = torch.randint(0, 2, (B,), device=device) * 2 - 1
+
+    versions = []
+    for ax in range(3):
+        for d in [-1, 1]:
+            s = torch.roll(out, shifts=d, dims=ax + 1)
+            if d == 1:
+                if ax == 0: s[:, 0, :, :] = 0
+                elif ax == 1: s[:, :, 0, :] = 0
+                else: s[:, :, :, 0] = 0
+            else:
+                if ax == 0: s[:, -1, :, :] = 0
+                elif ax == 1: s[:, :, -1, :] = 0
+                else: s[:, :, :, -1] = 0
+            versions.append(s)
+    versions.append(out)  # identity
+    stacked = torch.stack(versions, dim=0)  # (7, B, 8, 16, 16)
+
+    assign = torch.where(shift_sel, axes * 2 + (dirs == 1).long(), torch.full_like(axes, 6))
+    out = stacked[assign, torch.arange(B, device=device)]
+
+    return out
+
+
+# === Training =================================================================
+
+def train_vae_ca_classifier(model, train_loader, val_loader,
+                            n_epochs=60, lr=3e-3, weight_decay=1e-4,
+                            grad_clip=1.0, device='cuda',
+                            checkpoint_dir='/content/checkpoints_vae_ca'):
+    device = torch.device(device if torch.cuda.is_available() else 'cpu')
+    model = model.to(device)
+
+    amp_dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
+
+    optimizer = torch.optim.AdamW(model.parameters(), lr=lr, weight_decay=weight_decay)
+    scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=n_epochs, eta_min=lr * 0.01)
+
+    ce_loss_fn = nn.CrossEntropyLoss()
+    bce_loss_fn = nn.BCEWithLogitsLoss()
+
+    # From Cell 1 globals
+    dim_labels = torch.tensor([CLASS_META[n]["dim"] for n in CLASS_NAMES], dtype=torch.long, device=device)
+    curved_labels = torch.tensor([1.0 if CLASS_META[n]["curved"] else 0.0 for n in CLASS_NAMES], device=device)
+    curv_type_labels = torch.tensor([CURV_TO_IDX[CLASS_META[n]["curvature"]] for n in CLASS_NAMES], dtype=torch.long, device=device)
+
+    os.makedirs(checkpoint_dir, exist_ok=True)
+
+    best_acc = 0.0
+    print("=" * 70)
+
+    for epoch in range(1, n_epochs + 1):
+        model.train()
+        t0 = time.time()
+        total_loss = 0
+        correct = 0
+        total = 0
+
+        for grid, label in train_loader:
+            grid = grid.to(device, non_blocking=True)
+            label = label.to(device, non_blocking=True)
+
+            # Augmentation
+            grid = deform_grid(grid)
+
+            with torch.amp.autocast('cuda', dtype=amp_dtype):
+                out = model(grid, labels=label)
+
+                loss_cls = ce_loss_fn(out["class_logits"], label)
+                batch_dims = dim_labels[label]
+                loss_dim = ce_loss_fn(out["dim_logits"], batch_dims)
+                batch_curved = curved_labels[label].unsqueeze(-1)
+                loss_curved = bce_loss_fn(out["is_curved_pred"], batch_curved)
+                batch_curv_type = curv_type_labels[label]
+                loss_curv = ce_loss_fn(out["curv_type_logits"], batch_curv_type)
+
+                loss = loss_cls + 0.3 * loss_dim + 0.3 * loss_curved + 0.2 * loss_curv
+
+            optimizer.zero_grad(set_to_none=True)
+            if amp_dtype == torch.float16:
+                scaler = torch.amp.GradScaler('cuda')
+                scaler.scale(loss).backward()
+                scaler.unscale_(optimizer)
+                nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                loss.backward()
+                nn.utils.clip_grad_norm_(model.parameters(), grad_clip)
+                optimizer.step()
+
+            total_loss += loss.item()
+            pred = out["class_logits"].argmax(dim=-1)
+            correct += (pred == label).sum().item()
+            total += label.shape[0]
+
+        scheduler.step()
+
+        # === Validation ===
+        model.eval()
+        val_correct = 0
+        val_total = 0
+        curved_correct = 0
+        curved_total = 0
+        per_class_correct = torch.zeros(NUM_CLASSES, device=device)
+        per_class_total = torch.zeros(NUM_CLASSES, device=device)
+
+        with torch.no_grad():
+            for grid, label in val_loader:
+                grid = grid.to(device, non_blocking=True)
+                label = label.to(device, non_blocking=True)
+
+                with torch.amp.autocast('cuda', dtype=amp_dtype):
+                    out = model(grid)
+
+                pred = out["class_logits"].argmax(dim=-1)
+                val_correct += (pred == label).sum().item()
+                val_total += label.shape[0]
+
+                curved_pred = (out["is_curved_pred"].squeeze(-1) > 0.0).float()
+                curved_true = curved_labels[label]
+                curved_correct += (curved_pred == curved_true).sum().item()
+                curved_total += label.shape[0]
+
+                for c in range(NUM_CLASSES):
+                    mask = label == c
+                    per_class_total[c] += mask.sum()
+                    per_class_correct[c] += (pred[mask] == c).sum()
+
+        val_acc = val_correct / val_total
+        curved_acc = curved_correct / curved_total
+        train_acc = correct / total
+        train_loss = total_loss / len(train_loader)
+        elapsed = time.time() - t0
+        sps = total / elapsed
+
+        per_class_acc = per_class_correct / per_class_total.clamp(min=1)
+        worst = per_class_acc.argsort()[:10]
+
+        print(f"\nEpoch {epoch}/{n_epochs} | {elapsed:.1f}s | {sps:.0f} samp/s")
+        print(f"  Train: loss={train_loss:.4f} acc={train_acc:.4f}")
+        print(f"  Val:   acc={val_acc:.4f} curved={curved_acc:.4f}")
+        print(f"  LR: {scheduler.get_last_lr()[0]:.6f}")
+        print(f"  Worst classes:")
+        for idx in worst:
+            idx = idx.item()
+            acc = per_class_acc[idx].item() * 100
+            print(f"    {CLASS_NAMES[idx]:20s} {acc:5.1f}%")
+
+        if val_acc > best_acc:
+            best_acc = val_acc
+            torch.save(model.state_dict(), os.path.join(checkpoint_dir, 'best.pt'))
+            # Also save to flat path for Cell 5 compat
+            torch.save(model.state_dict(), '/content/best_vae_ca_classifier.pt')
+            print(f"  ★ New best: {best_acc:.2%}")
+
+        # Always save latest
+        torch.save(model.state_dict(), os.path.join(checkpoint_dir, 'latest.pt'))
+
+    print(f"\n{'=' * 70}")
+    print(f"Training complete. Best val acc: {best_acc:.2%}")
+    return best_acc
+
+
+# === Run ======================================================================
+
+print("=" * 70)
+print(f"Patch Cross-Attention Shape Classifier — {GZ}×{GY}×{GX}")
+print(f"  Patches: {MACRO_Z}×{MACRO_Y}×{MACRO_X} = {MACRO_N} of {PATCH_Z}×{PATCH_Y}×{PATCH_X}")
+print("=" * 70)
+
+print("\nGenerating training data...")
+gen = ShapeGenerator(seed=42)
+train_data = gen.generate_dataset(2000, seed=42)
+val_data = gen.generate_dataset(400, seed=999)
+
+print(f"  Generated {len(train_data['grids'])} train + {len(val_data['grids'])} val")
+print(f"  Classes: {NUM_CLASSES}")
+occ = train_data['grids'].reshape(len(train_data['grids']), -1).sum(axis=1)
+print(f"  Avg occupied voxels: {occ.mean():.1f} / {GZ*GY*GX} ({occ.mean()/(GZ*GY*GX)*100:.1f}%)")
+
+batch_size = 1024
+train_ds = TensorDataset(torch.from_numpy(train_data['grids']).float(),
+                         torch.from_numpy(train_data['labels']).long())
+val_ds = TensorDataset(torch.from_numpy(val_data['grids']).float(),
+                       torch.from_numpy(val_data['labels']).long())
+
+train_loader = DataLoader(train_ds, batch_size=batch_size, shuffle=True,
+                          num_workers=2, pin_memory=True, drop_last=True)
+val_loader = DataLoader(val_ds, batch_size=batch_size, shuffle=False,
+                        num_workers=2, pin_memory=True)
+
+print(f"  Batch size: {batch_size}")
+print(f"  Train batches: {len(train_loader)} | Val batches: {len(val_loader)}")
+
+model = PatchCrossAttentionClassifier(n_classes=NUM_CLASSES)
+n_params = sum(p.numel() for p in model.parameters())
+print(f"  Model: {n_params:,} params")
+
+device = 'cuda' if torch.cuda.is_available() else 'cpu'
+print(f"  Device: {device}")
+
+best_acc = train_vae_ca_classifier(
+    model, train_loader, val_loader,
+    n_epochs=60, lr=3e-3, weight_decay=1e-4, device=device)
+
+print(f"\nDone. Best accuracy: {best_acc:.2%}")