Add train.py

train.py

@@ -0,0 +1,375 @@
+#!/usr/bin/env python3
+"""
+Train Venice-H1 re-ranker on pre-extracted DeRIS feature caches.
+
+Matches paper Section 3.5 exactly:
+  - Loss: L = L_gate + λ·L_gain (λ=5)
+  - L_gate: focal BCE (γ=2, auto w_pos)
+  - L_gain: smooth-L1 IoU regression on ALL samples
+  - AdamW (lr=5e-4, wd=1e-4), cosine + 3-epoch warmup
+  - 20 epochs, batch 512, FP16
+  - Training time: ~3 min on single RTX 3090
+
+Usage:
+    python train.py --config venice_h1/configs/default.yaml
+    python train.py --config venice_h1/configs/default.yaml --no_grid
+"""
+
+import argparse
+import os
+import random
+from pathlib import Path
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.data import DataLoader, Dataset, ConcatDataset
+from torch.cuda.amp import GradScaler, autocast
+
+try:
+    from torch.utils.tensorboard import SummaryWriter
+    _HAS_TB = True
+except ImportError:
+    _HAS_TB = False
+
+try:
+    import yaml
+except ImportError:
+    raise SystemExit("pip install pyyaml")
+
+try:
+    from sklearn.metrics import roc_auc_score
+    _HAS_SKLEARN = True
+except ImportError:
+    _HAS_SKLEARN = False
+
+from venice_h1.model.reranker import VeniceH1Reranker
+
+
+def set_seed(seed: int):
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+
+
+def load_config(path: str) -> dict:
+    with open(path) as f:
+        return yaml.safe_load(f)
+
+
+# ---------------------------------------------------------------------------
+# Dataset
+# ---------------------------------------------------------------------------
+
+class FeatureCacheDataset(Dataset):
+    """Loads a pre-extracted .pt feature cache from extract_features.py."""
+
+    def __init__(self, path: str, use_grid: bool = True):
+        data = torch.load(path, map_location="cpu")
+        self.use_grid = use_grid
+
+        self.query_feat  = data["query_feat"].float()
+        self.det_scores  = data["det_scores"].float()
+        self.query_ious  = data["query_ious"].float()
+        self.oracle_idx  = data["oracle_idx"].long()
+
+        self.mask_mean   = data["mask_mean"].float()
+        self.mask_max    = data["mask_max"].float()
+        self.mask_area   = data["mask_area"].float()
+        self.mask_std    = data["mask_std"].float()
+
+        if use_grid:
+            self.grid_mean_4  = data["grid_mean_4"].float()
+            self.grid_max_4   = data["grid_max_4"].float()
+            self.boundary_4   = data["boundary_4"].float()
+            self.grid_mean_8  = data["grid_mean_8"].float()
+            self.grid_max_8   = data["grid_max_8"].float()
+            self.boundary_8   = data["boundary_8"].float()
+            self.grid_mean_16 = data["grid_mean_16"].float()
+            self.grid_max_16  = data["grid_max_16"].float()
+            self.boundary_16  = data["boundary_16"].float()
+
+        self.failure_flag = (self.oracle_idx != 0).float()
+
+    def __len__(self):
+        return len(self.oracle_idx)
+
+    def __getitem__(self, idx):
+        N = self.query_feat.shape[1]
+        qf = self.query_feat[idx]                         # [N, D]
+        ds = self.det_scores[idx].unsqueeze(-1)           # [N, 1]
+        mm = self.mask_mean[idx].unsqueeze(-1)            # [N, 1]
+        mx = self.mask_max[idx].unsqueeze(-1)             # [N, 1]
+        ma = self.mask_area[idx].unsqueeze(-1)            # [N, 1]
+        ms = self.mask_std[idx].unsqueeze(-1)             # [N, 1]
+
+        parts = [qf, ds, mm, mx, ma, ms]                 # [N, D+5]
+
+        if self.use_grid:
+            gm4 = self.grid_mean_4[idx]                   # [N, 16]
+            gx4 = self.grid_max_4[idx]                    # [N, 16]
+            b4  = self.boundary_4[idx].unsqueeze(-1)      # [N, 1]
+            gm8 = self.grid_mean_8[idx]                   # [N, 64]
+            gx8 = self.grid_max_8[idx]                    # [N, 64]
+            b8  = self.boundary_8[idx].unsqueeze(-1)      # [N, 1]
+            gm16 = self.grid_mean_16[idx]                 # [N, 256]
+            gx16 = self.grid_max_16[idx]                  # [N, 256]
+            b16  = self.boundary_16[idx].unsqueeze(-1)    # [N, 1]
+            parts += [gm4, gx4, b4, gm8, gx8, b8, gm16, gx16, b16]
+
+        features = torch.cat(parts, dim=-1)               # [N, 936]
+
+        return {
+            "features":     features,
+            "det_scores":   self.det_scores[idx],
+            "mask_means":   self.mask_mean[idx],
+            "oracle_idx":   self.oracle_idx[idx],
+            "failure_flag": self.failure_flag[idx],
+            "query_ious":   self.query_ious[idx],
+        }
+
+
+# ---------------------------------------------------------------------------
+# Loss (Section 3.5)
+# ---------------------------------------------------------------------------
+
+def focal_bce_loss(pred: torch.Tensor, target: torch.Tensor,
+                   gamma: float = 2.0, pos_weight: float | None = None):
+    """Focal binary cross-entropy (Eq. in Section 3.5)."""
+    if pos_weight is not None:
+        weight = torch.where(target > 0.5, pos_weight, 1.0)
+    else:
+        weight = None
+
+    bce = F.binary_cross_entropy(pred, target, reduction='none')
+    pt = pred * target + (1 - pred) * (1 - target)
+    focal = ((1 - pt) ** gamma) * bce
+
+    if weight is not None:
+        focal = focal * weight
+    return focal.mean()
+
+
+def compute_loss(out: dict, batch: dict, cfg: dict) -> dict:
+    """L = L_gate + λ·L_gain (Section 3.5)."""
+    device = out["p_fail"].device
+    tcfg = cfg["training"]
+
+    failure_flag = batch["failure_flag"].to(device)
+    query_ious   = batch["query_ious"].to(device)
+
+    # Auto positive weight for focal BCE
+    n_pos = failure_flag.sum().clamp(min=1)
+    n_neg = (1 - failure_flag).sum().clamp(min=1)
+    wpos = (n_neg / n_pos).clamp(max=50.0) if tcfg.get("auto_wpos") else None
+
+    # L_gate: focal BCE
+    loss_gate = focal_bce_loss(
+        out["p_fail"], failure_flag,
+        gamma=tcfg["focal_gamma"], pos_weight=wpos)
+
+    # L_gain: smooth-L1 on IoU gain (all samples, dense supervision)
+    gain_target = query_ious - query_ious[:, 0:1]          # relative to Q0
+    loss_gain = F.smooth_l1_loss(out["gain_logits"], gain_target)
+
+    # Total loss
+    total = loss_gate + tcfg["lambda_gain"] * loss_gain
+
+    return {"total": total, "gate": loss_gate, "gain": loss_gain}
+
+
+# ---------------------------------------------------------------------------
+# Evaluation
+# ---------------------------------------------------------------------------
+
+@torch.no_grad()
+def evaluate(model, loader, device, tau: float = 0.05) -> dict:
+    model.eval()
+    total = failures = reranked_count = correct = harmful = 0
+    all_pfail, all_flag = [], []
+
+    for batch in loader:
+        features     = batch["features"].to(device)
+        det_scores   = batch["det_scores"].to(device)
+        mask_means   = batch["mask_means"].to(device)
+        oracle_idx   = batch["oracle_idx"].to(device)
+        failure_flag = batch["failure_flag"].to(device)
+        query_ious   = batch["query_ious"].to(device)
+
+        out = model(features, det_scores, mask_means)
+        selected = model.rerank(features, det_scores, mask_means, tau=tau)
+
+        B = len(oracle_idx)
+        total += B
+        fail_mask = failure_flag.bool()
+        failures += fail_mask.sum().item()
+
+        reranked_mask = selected != 0
+        reranked_count += reranked_mask.sum().item()
+        correct += (selected[fail_mask] == oracle_idx[fail_mask]).sum().item()
+
+        q0_iou  = query_ious[:, 0]
+        sel_iou = query_ious.gather(1, selected.unsqueeze(1)).squeeze(1)
+        harmful += (reranked_mask & (sel_iou < q0_iou - 0.01)).sum().item()
+
+        all_pfail.append(out["p_fail"].cpu())
+        all_flag.append(failure_flag.cpu())
+
+    all_pfail = torch.cat(all_pfail).numpy()
+    all_flag  = torch.cat(all_flag).numpy()
+    auc = roc_auc_score(all_flag, all_pfail) if _HAS_SKLEARN and all_flag.sum() > 0 else 0.0
+
+    return {
+        "failure_rate":    failures / max(total, 1),
+        "rerank_rate":     reranked_count / max(total, 1),
+        "correct_rerank":  correct / max(failures, 1),
+        "harmful_switch":  harmful / max(total, 1),
+        "gate_auc":        auc,
+    }
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+def main():
+    parser = argparse.ArgumentParser(description="Train Venice-H1 re-ranker")
+    parser.add_argument("--config", default="venice_h1/configs/default.yaml")
+    parser.add_argument("--no_grid", action="store_true",
+                        help="Ablation: BASE features only (Df=261)")
+    parser.add_argument("--tau", type=float, default=None)
+    args = parser.parse_args()
+
+    cfg = load_config(args.config)
+    if args.no_grid:
+        cfg["model"]["use_grid"] = False
+
+    set_seed(cfg["training"]["seed"])
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    print(f"Device: {device}")
+
+    use_grid = cfg["model"]["use_grid"]
+    tau = args.tau if args.tau is not None else cfg["model"]["tau"]
+
+    # ---- Data ----
+    train_paths = [p for p in cfg["data"]["train_splits"] if Path(p).exists()]
+    val_paths   = [p for p in cfg["data"]["val_splits"] if Path(p).exists()]
+
+    if not train_paths:
+        print("ERROR: No training data found. Run extract_features.py first.")
+        print("Expected paths:", cfg["data"]["train_splits"])
+        return
+
+    train_ds = ConcatDataset([FeatureCacheDataset(p, use_grid) for p in train_paths])
+    val_ds   = ConcatDataset([FeatureCacheDataset(p, use_grid) for p in val_paths]) \
+               if val_paths else None
+
+    bs = cfg["training"]["batch_size"]
+    train_loader = DataLoader(train_ds, batch_size=bs, shuffle=True,
+                              num_workers=4, pin_memory=True, drop_last=True)
+    val_loader = DataLoader(val_ds, batch_size=bs, shuffle=False,
+                            num_workers=4, pin_memory=True) if val_ds else None
+
+    print(f"Train: {len(train_ds):,} samples  |  "
+          f"Val: {len(val_ds) if val_ds else 0:,} samples")
+    print(f"Feature set: {'BASE+GRID (Df=936)' if use_grid else 'BASE (Df=261)'}")
+
+    # ---- Model ----
+    model = VeniceH1Reranker(**cfg["model"]).to(device)
+    print(f"Venice-H1: {model.num_parameters():,} parameters (~11.3M)")
+
+    # ---- Optimizer (Section 3.5: AdamW, lr=5e-4, wd=1e-4) ----
+    optimizer = torch.optim.AdamW(
+        model.parameters(),
+        lr=cfg["training"]["lr"],
+        weight_decay=cfg["training"]["weight_decay"])
+
+    epochs = cfg["training"]["epochs"]
+    warmup = cfg["training"]["warmup_epochs"]
+
+    def lr_lambda(epoch):
+        if epoch < warmup:
+            return epoch / warmup
+        progress = (epoch - warmup) / (epochs - warmup)
+        return 0.5 * (1 + np.cos(np.pi * progress))
+
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda)
+
+    # FP16
+    scaler = GradScaler() if cfg["training"].get("fp16") else None
+
+    # Logging
+    writer = SummaryWriter(cfg["output"]["log_dir"]) if _HAS_TB else None
+    ckpt_dir = cfg["output"]["checkpoint_dir"]
+    os.makedirs(ckpt_dir, exist_ok=True)
+
+    best_auc = 0.0
+
+    # ---- Training Loop ----
+    print(f"\nTraining for {epochs} epochs (batch={bs}, lr={cfg['training']['lr']})")
+    print("-" * 70)
+
+    for epoch in range(1, epochs + 1):
+        model.train()
+        epoch_loss = 0.0
+
+        for batch in train_loader:
+            features   = batch["features"].to(device)
+            det_scores = batch["det_scores"].to(device)
+            mask_means = batch["mask_means"].to(device)
+
+            optimizer.zero_grad()
+
+            if scaler:
+                with autocast():
+                    out = model(features, det_scores, mask_means)
+                    losses = compute_loss(out, batch, cfg)
+                scaler.scale(losses["total"]).backward()
+                scaler.unscale_(optimizer)
+                nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                scaler.step(optimizer)
+                scaler.update()
+            else:
+                out = model(features, det_scores, mask_means)
+                losses = compute_loss(out, batch, cfg)
+                losses["total"].backward()
+                nn.utils.clip_grad_norm_(model.parameters(), 1.0)
+                optimizer.step()
+
+            epoch_loss += losses["total"].item()
+
+        scheduler.step()
+        avg_loss = epoch_loss / len(train_loader)
+
+        # Validation
+        metrics = evaluate(model, val_loader, device, tau=tau) if val_loader else {}
+
+        if writer:
+            writer.add_scalar("train/loss", avg_loss, epoch)
+            for k, v in metrics.items():
+                writer.add_scalar(f"val/{k}", v, epoch)
+
+        auc = metrics.get("gate_auc", 0)
+        harm = metrics.get("harmful_switch", 0) * 100
+        print(f"  Epoch {epoch:2d}/{epochs} | loss={avg_loss:.4f} | "
+              f"AUC={auc:.3f} | harmful={harm:.2f}%")
+
+        if auc > best_auc:
+            best_auc = auc
+            torch.save(
+                {"epoch": epoch, "model": model.state_dict(),
+                 "config": cfg, "metrics": metrics},
+                os.path.join(ckpt_dir, "best.pt"))
+
+    if writer:
+        writer.close()
+
+    print(f"\nDone. Best Gate AUC: {best_auc:.4f}")
+    print(f"Checkpoint: {os.path.join(ckpt_dir, 'best.pt')}")
+
+
+if __name__ == "__main__":
+    main()