#!/usr/bin/env python3 """ Train a NexarTemporalHead (or NexarSimpleHead) on domain-adapted Nexar features. The SFT backbone is already frozen and features are pre-cached. This trainer only optimises the lightweight collision prediction head. Usage: python -m training.Nexar.nexar_trainer \ --cache_pos data/nexar_cache/train_positive.pt \ --cache_neg data/nexar_cache/train_negative.pt \ --output_dir checkpoints/Nexar/nexar_v1 \ --arch temporal \ --n_windows 3 \ --epochs 20 \ --batch_size 64 \ --lr 3e-4 """ from __future__ import annotations import argparse import json import logging import random from pathlib import Path from typing import List, Optional, Tuple import numpy as np import torch import torch.nn as nn import torch.nn.functional as F from torch.optim import AdamW from torch.optim.lr_scheduler import CosineAnnealingLR from torch.utils.data import DataLoader, WeightedRandomSampler from sklearn.metrics import average_precision_score, roc_auc_score import sys sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from training.Nexar.nexar_dataset import NexarTrainDataset, nexar_collate_train from training.Nexar.nexar_model import build_model logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Nexar.trainer") SEED = 42 def set_seed(seed: int): random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) torch.cuda.manual_seed_all(seed) def split_dataset(ds: NexarTrainDataset, val_frac: float = 0.15, seed: int = SEED): """Stratified train/val split.""" labels = [s["label"] for s in ds.samples] pos_idx = [i for i, l in enumerate(labels) if l == 1] neg_idx = [i for i, l in enumerate(labels) if l == 0] rng = random.Random(seed) rng.shuffle(pos_idx) rng.shuffle(neg_idx) n_val_pos = max(1, int(len(pos_idx) * val_frac)) n_val_neg = max(1, int(len(neg_idx) * val_frac)) val_idx = pos_idx[:n_val_pos] + neg_idx[:n_val_neg] train_idx = pos_idx[n_val_pos:] + neg_idx[n_val_neg:] from torch.utils.data import Subset return Subset(ds, train_idx), Subset(ds, val_idx) def make_sampler(subset) -> WeightedRandomSampler: """Class-balanced weighted sampler for the training subset.""" labels = [subset.dataset.samples[i]["label"] for i in subset.indices] labels_arr = np.array(labels, dtype=float) n_pos = labels_arr.sum() n_neg = len(labels_arr) - n_pos weights = np.where(labels_arr == 1, len(labels_arr) / (2 * n_pos + 1e-9), len(labels_arr) / (2 * n_neg + 1e-9)) return WeightedRandomSampler( weights=torch.from_numpy(weights).float(), num_samples=len(subset), replacement=True, ) def compute_ap(labels: np.ndarray, scores: np.ndarray) -> float: try: return float(average_precision_score(labels, scores)) except Exception: return float("nan") def train_epoch(model, loader, optimizer, device) -> float: model.train() total_loss = 0.0 n = 0 for batch in loader: beliefs = batch["beliefs"].to(device) # [B, T, H] or [B, H] tta_means = batch["tta_means"].to(device) tta_vars = batch["tta_vars"].to(device) p_alerts = batch["p_alerts"].to(device) labels = batch["labels"].to(device) # [B] float if isinstance(model, torch.nn.Module) and hasattr(model, "lstm"): # Temporal model: [B, T, H] scores = model(beliefs, tta_means, tta_vars, p_alerts) else: # Simple model: last window only [B, H] scores = model(beliefs[:, -1, :], tta_means[:, -1], tta_vars[:, -1], p_alerts[:, -1]) loss = F.binary_cross_entropy(scores, labels) optimizer.zero_grad() loss.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) optimizer.step() total_loss += loss.item() * len(labels) n += len(labels) return total_loss / max(n, 1) @torch.no_grad() def eval_epoch(model, loader, device) -> Tuple[float, float, float]: model.eval() all_scores: List[float] = [] all_labels: List[float] = [] total_loss = 0.0 n = 0 for batch in loader: beliefs = batch["beliefs"].to(device) tta_means = batch["tta_means"].to(device) tta_vars = batch["tta_vars"].to(device) p_alerts = batch["p_alerts"].to(device) labels = batch["labels"].to(device) if hasattr(model, "lstm"): scores = model(beliefs, tta_means, tta_vars, p_alerts) else: scores = model(beliefs[:, -1, :], tta_means[:, -1], tta_vars[:, -1], p_alerts[:, -1]) loss = F.binary_cross_entropy(scores, labels) total_loss += loss.item() * len(labels) n += len(labels) all_scores.extend(scores.cpu().tolist()) all_labels.extend(labels.cpu().tolist()) arr_l = np.array(all_labels) arr_s = np.array(all_scores) ap = compute_ap(arr_l, arr_s) try: auc = float(roc_auc_score(arr_l, arr_s)) except Exception: auc = float("nan") return total_loss / max(n, 1), ap, auc def main(): parser = argparse.ArgumentParser("nexar_trainer") parser.add_argument("--cache_pos", required=True, help=".pt cache for positive train videos") parser.add_argument("--cache_neg", required=True, help=".pt cache for negative train videos") parser.add_argument("--output_dir", required=True) parser.add_argument("--arch", default="temporal", choices=["simple", "temporal"]) parser.add_argument("--n_windows", type=int, default=3) parser.add_argument("--epochs", type=int, default=30) parser.add_argument("--batch_size", type=int, default=64) parser.add_argument("--lr", type=float, default=3e-4) parser.add_argument("--lr_min", type=float, default=1e-6) parser.add_argument("--weight_decay",type=float, default=1e-4) parser.add_argument("--val_frac", type=float, default=0.15) parser.add_argument("--patience", type=int, default=8) parser.add_argument("--hidden_dim", type=int, default=2048, help="SFT hidden_dim (Qwen2.5-VL-3B = 2048)") args = parser.parse_args() set_seed(SEED) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") out_dir = Path(args.output_dir) out_dir.mkdir(parents=True, exist_ok=True) # ── data ───────────────────────────────────────────────────────────────── full_ds = NexarTrainDataset(args.cache_pos, args.cache_neg, n_windows=args.n_windows) train_subset, val_subset = split_dataset(full_ds, val_frac=args.val_frac) logger.info(f"Train: {len(train_subset)} Val: {len(val_subset)}") sampler = make_sampler(train_subset) train_loader = DataLoader(train_subset, batch_size=args.batch_size, sampler=sampler, num_workers=4, collate_fn=nexar_collate_train, pin_memory=True) val_loader = DataLoader(val_subset, batch_size=args.batch_size, shuffle=False, num_workers=4, collate_fn=nexar_collate_train, pin_memory=True) # ── model ───────────────────────────────────────────────────────────────── model = build_model(args.hidden_dim, args.arch).to(device) total_params = sum(p.numel() for p in model.parameters()) logger.info(f"NexarHead ({args.arch}): {total_params:,} params") optimizer = AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) total_steps = args.epochs * len(train_loader) scheduler = CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=args.lr_min) # ── training loop ───────────────────────────────────────────────────────── best_ap = 0.0 patience_count = 0 history = [] for epoch in range(1, args.epochs + 1): train_loss = train_epoch(model, train_loader, optimizer, device) scheduler.step() val_loss, val_ap, val_auc = eval_epoch(model, val_loader, device) lr = optimizer.param_groups[0]["lr"] logger.info( f"Epoch {epoch:3d}/{args.epochs} " f"train_loss={train_loss:.4f} val_loss={val_loss:.4f} " f"val_AP={val_ap:.4f} val_AUC={val_auc:.4f} lr={lr:.2e}" ) history.append({ "epoch": epoch, "train_loss": train_loss, "val_loss": val_loss, "val_ap": val_ap, "val_auc": val_auc, }) if val_ap > best_ap: best_ap = val_ap patience_count = 0 torch.save(model.state_dict(), out_dir / "best_model.pt") with open(out_dir / "best_meta.json", "w") as f: json.dump({"epoch": epoch, "val_ap": val_ap, "val_auc": val_auc, "arch": args.arch, "hidden_dim": args.hidden_dim, "n_windows": args.n_windows}, f, indent=2) logger.info(f" ★ New best val_AP={best_ap:.4f} — checkpoint saved") else: patience_count += 1 if patience_count >= args.patience: logger.info(f"Early stopping at epoch {epoch} (patience={args.patience})") break with open(out_dir / "history.json", "w") as f: json.dump(history, f, indent=2) logger.info(f"\n✅ Training complete. Best val_AP = {best_ap:.4f}") logger.info(f" Checkpoint: {out_dir}/best_model.pt") if __name__ == "__main__": main()