training/Nexar/train_resnet18_bigru.py

#!/usr/bin/env python3
"""ResNet18-BiGRU Track-A trainer.

Protocol (CORRECTED 2026-04-28):
  - TRAIN  = all 1500 nexar train clips (data/rgb_clip_cache/train.pt)
  - VAL    = test-public subset (~672 clips) of data/rgb_clip_cache/test.pt,
              filtered by NEXAR_COLLISION/solution.csv `Usage == "Public"`.
  - TEST   = test-private subset (~672 clips, `Usage == "Private"`).

This replaces the earlier 1280/220 internal-val split.

Cache (`tools/extract_rgb_clip_cache.py`) contains T=64 [N, T, 512] feats
(frozen-ResNet18 ImageNet pooled features).

Model: BiGRU(d=256) → attentive pool → MLP → 1 logit.

Usage:
    python -m training.Nexar.train_resnet18_bigru --seed 0 \\
      --train_cache data/rgb_clip_cache/train.pt \\
      --test_cache  data/rgb_clip_cache/test.pt \\
      --solution_csv NEXAR_COLLISION/solution.csv \\
      --label_csv    nexar-collision-prediction/train.csv \\
      --output       checkpoints/Nexar/resnet18_bigru_seed0
"""
from __future__ import annotations

import argparse
import csv
import json
import random
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn
from sklearn.metrics import average_precision_score, roc_auc_score
from torch.utils.data import DataLoader, Dataset


def set_seed(s: int):
    random.seed(s); np.random.seed(s); torch.manual_seed(s)
    torch.cuda.manual_seed_all(s)


class FeatDataset(Dataset):
    def __init__(self, cache_path: Path, labels: dict[str, int],
                 badas_soft: dict | None = None,
                 keep_ids: set[str] | None = None):
        c = torch.load(cache_path, weights_only=False, map_location="cpu")
        self.ids:  list[str] = c["ids"]
        self.feat = c["feat"]
        self.labels = labels
        self.badas_soft = badas_soft or {}
        kept = []
        for i, vid in enumerate(self.ids):
            if vid not in labels:
                continue
            if keep_ids is not None and vid not in keep_ids:
                continue
            kept.append(i)
        self.kept = kept

    def __len__(self):
        return len(self.kept)

    def __getitem__(self, i):
        idx = self.kept[i]
        vid = self.ids[idx]
        x = self.feat[idx].float()
        y = float(self.labels[vid])
        soft = float(self.badas_soft.get(vid, y))
        return x, torch.tensor(y), torch.tensor(soft), vid


class BiGRUHead(nn.Module):
    def __init__(self, in_dim: int = 512, hidden: int = 256,
                 dropout: float = 0.3):
        super().__init__()
        self.gru = nn.GRU(in_dim, hidden, num_layers=1, bidirectional=True,
                           batch_first=True)
        self.attn = nn.Sequential(nn.Linear(2 * hidden, 128),
                                    nn.Tanh(), nn.Linear(128, 1))
        self.cls  = nn.Sequential(nn.Dropout(dropout),
                                    nn.Linear(2 * hidden, 1))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        h, _ = self.gru(x)
        a = torch.softmax(self.attn(h).squeeze(-1), dim=1)
        pooled = (h * a.unsqueeze(-1)).sum(dim=1)
        return self.cls(pooled).squeeze(-1)


def load_labels(csv_path: Path) -> dict[str, int]:
    return {r["id"]: int(r["target"] or 0)
            for r in csv.DictReader(open(csv_path))
            if r.get("target") is not None}


def load_solution_split(csv_path: Path):
    rows = list(csv.DictReader(open(csv_path)))
    pub  = {r["id"] for r in rows if r["Usage"] == "Public"}
    priv = {r["id"] for r in rows if r["Usage"] == "Private"}
    targets = {r["id"]: int(r["target"]) for r in rows}
    return pub, priv, targets


def load_badas_soft(p: Path) -> dict[str, float]:
    if not p.exists():
        return {}
    j = json.loads(p.read_text())
    return {cid: float(rec["score_last4s"])
            for cid, rec in j.items()
            if rec.get("score_last4s") is not None
            and not np.isnan(rec["score_last4s"])}


@torch.no_grad()
def eval_split(model, loader, device) -> tuple[float, float]:
    model.eval()
    ys, ps = [], []
    for x, y, _soft, _vid in loader:
        logit = model(x.to(device)).cpu().numpy()
        ys.extend(y.numpy().tolist()); ps.extend(logit.tolist())
    ys, ps = np.asarray(ys), np.asarray(ps)
    if len(np.unique(ys)) < 2:
        return float("nan"), float("nan")
    return (float(average_precision_score(ys, ps)),
            float(roc_auc_score(ys, ps)))


@torch.no_grad()
def eval_kaggle_mAP(model, ds: FeatDataset, device, solution: Path,
                     batch: int = 64) -> tuple[float, float]:
    rows = list(csv.DictReader(open(solution)))
    group = {r["id"]: int(r["group"]) for r in rows}
    usage = {r["id"]: r["Usage"] for r in rows}
    model.eval()
    score: dict[str, float] = {}
    target: dict[str, int] = {}
    dl = DataLoader(ds, batch_size=batch, shuffle=False, num_workers=4)
    for x, y, _soft, vids in dl:
        logit = model(x.to(device)).cpu().numpy()
        for v, p, t in zip(vids, logit.tolist(), y.numpy().tolist()):
            score[v] = float(p); target[v] = int(t)
    common = sorted(set(score) & set(group))
    if not common:
        return float("nan"), float("nan")
    pub_aps, priv_aps = [], []
    for g in (0, 1, 2):
        for u, sink in (("Public", pub_aps), ("Private", priv_aps)):
            ids = [v for v in common if usage.get(v) == u and group.get(v) == g]
            if len(ids) < 2: continue
            y = np.array([target[v] for v in ids])
            s = np.array([score[v] for v in ids])
            if len(np.unique(y)) < 2: continue
            sink.append(float(average_precision_score(y, s)))
    pub_mAP  = float(np.mean(pub_aps))  if pub_aps  else float("nan")
    priv_mAP = float(np.mean(priv_aps)) if priv_aps else float("nan")
    return pub_mAP, priv_mAP


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--seed",        type=int, default=0)
    ap.add_argument("--train_cache", default="data/rgb_clip_cache/train.pt")
    ap.add_argument("--test_cache",  default="data/rgb_clip_cache/test.pt")
    ap.add_argument("--solution_csv", default="NEXAR_COLLISION/solution.csv")
    ap.add_argument("--label_csv",   default="nexar-collision-prediction/train.csv")
    ap.add_argument("--output",      default="checkpoints/Nexar/resnet18_bigru_seed0")
    ap.add_argument("--epochs",   type=int,   default=25)
    ap.add_argument("--batch",    type=int,   default=64)
    ap.add_argument("--lr",       type=float, default=3e-4)
    ap.add_argument("--wd",       type=float, default=1e-4)
    ap.add_argument("--hidden",   type=int,   default=256)
    ap.add_argument("--dropout",  type=float, default=0.3)
    ap.add_argument("--badas_soft_label", default=None)
    ap.add_argument("--alpha_soft", type=float, default=0.3)
    ap.add_argument("--report_private", action="store_true")
    args = ap.parse_args()

    set_seed(args.seed)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    out = Path(args.output); out.mkdir(parents=True, exist_ok=True)

    badas_soft = (load_badas_soft(Path(args.badas_soft_label))
                   if args.badas_soft_label else {})
    if badas_soft:
        print(f"[init] loaded {len(badas_soft)} BADAS soft labels")

    train_labels = load_labels(Path(args.label_csv))
    print(f"[init] train labels: {len(train_labels)}")

    pub_ids, priv_ids, test_targets = load_solution_split(Path(args.solution_csv))
    print(f"[init] solution split: public={len(pub_ids)}  private={len(priv_ids)}")

    train_ds = FeatDataset(Path(args.train_cache), train_labels, badas_soft)
    val_ds   = FeatDataset(Path(args.test_cache), test_targets, keep_ids=pub_ids)
    priv_ds  = FeatDataset(Path(args.test_cache), test_targets, keep_ids=priv_ids)
    print(f"[init] datasets: train_n={len(train_ds)}  val_pub_n={len(val_ds)}  "
          f"priv_n={len(priv_ds)}")

    train_dl = DataLoader(train_ds, batch_size=args.batch, shuffle=True,
                            num_workers=4, pin_memory=True)
    val_dl   = DataLoader(val_ds,   batch_size=args.batch, shuffle=False,
                            num_workers=4, pin_memory=True)

    model = BiGRUHead(in_dim=512, hidden=args.hidden,
                       dropout=args.dropout).to(device)
    opt   = torch.optim.AdamW(model.parameters(), lr=args.lr,
                               weight_decay=args.wd)
    sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=args.epochs)
    bce   = nn.BCEWithLogitsLoss()

    best_pub_mAP = -1.0; best_path = out / "best.pt"
    history = []
    for ep in range(args.epochs):
        model.train(); running = 0.0; n_batch = 0
        for x, y, soft, _vid in train_dl:
            x = x.to(device); y = y.to(device); soft = soft.to(device)
            logit = model(x)
            loss_gt = bce(logit, y)
            loss = ((1 - args.alpha_soft) * loss_gt +
                     args.alpha_soft * bce(logit, soft)) if badas_soft else loss_gt
            opt.zero_grad(); loss.backward(); opt.step()
            running += float(loss.item()); n_batch += 1
        sched.step()
        pub_ap, pub_auc = eval_split(model, val_dl, device)
        pub_mAP, _ = eval_kaggle_mAP(model, val_ds, device,
                                        Path(args.solution_csv),
                                        batch=args.batch)
        priv_str = ""
        if args.report_private:
            priv_dl = DataLoader(priv_ds, batch_size=args.batch, shuffle=False,
                                   num_workers=4, pin_memory=True)
            priv_ap, _ = eval_split(model, priv_dl, device)
            _, priv_mAP = eval_kaggle_mAP(model, priv_ds, device,
                                            Path(args.solution_csv),
                                            batch=args.batch)
            priv_str = f"  [info] priv_mAP={priv_mAP:.4f}"
        avg = running / max(n_batch, 1)
        print(f"epoch {ep+1:02d}/{args.epochs}  loss={avg:.4f}  "
              f"pub_AP={pub_ap:.4f}  pub_AUC={pub_auc:.4f}  "
              f"pub_mAP={pub_mAP:.4f}{priv_str}", flush=True)
        history.append({"epoch": ep + 1, "loss": avg,
                          "pub_AP": pub_ap, "pub_AUC": pub_auc,
                          "pub_mAP": pub_mAP})
        if pub_mAP > best_pub_mAP:
            best_pub_mAP = pub_mAP
            torch.save({"head_state": model.state_dict(),
                          "args": vars(args), "epoch": ep + 1,
                          "pub_AP": pub_ap, "pub_AUC": pub_auc,
                          "pub_mAP": pub_mAP}, best_path)

    (out / "history.json").write_text(json.dumps(history, indent=2))
    print(f"[done] best pub_mAP={best_pub_mAP:.4f}  ckpt={best_path}")


if __name__ == "__main__":
    main()