train.py

#!/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()