"""Phase G.0c — Re-train DangerHead with 8-way hazard auxiliary head. Joint loss = existing alert-binary BCE (per-frame + clip) + 0.3 · CE(hazard_logits, hazard_target) Hazard targets come from `data/policy_labels/hazard_categories_*.json` built by `tools/build_hazard_labels.py`. Output: checkpoints/danger_v3_hazard/best.pt """ from __future__ import annotations import argparse import json import logging import sys from pathlib import Path import torch import torch.nn.functional as F from torch.utils.data import DataLoader, Dataset from tqdm import tqdm from sklearn.metrics import (accuracy_score, balanced_accuracy_score, average_precision_score, roc_auc_score) import numpy as np ROOT = Path(__file__).resolve().parents[2] sys.path.insert(0, str(ROOT)) from lkalert.models.danger_head import DangerHead, danger_loss logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("danger_hazard") N_HAZARDS = 8 class HazardDataset(Dataset): def __init__(self, cache_path: Path, hazard_path: Path): self.cache = torch.load(cache_path, weights_only=False, map_location="cpu") hz = json.loads(hazard_path.read_text()) # Index hazard labels by video_id (parallel to cache['ids']) ids_to_h = dict(zip(hz["ids"], hz["labels"])) self.hazard = torch.tensor( [ids_to_h.get(vid, 7) for vid in self.cache["ids"]], dtype=torch.long) logger.info(f" loaded {cache_path.name}: N={len(self.cache['ids'])} " f"hazard dist={torch.bincount(self.hazard, minlength=N_HAZARDS).tolist()}") def __len__(self): return len(self.cache["ids"]) def __getitem__(self, idx): return { "belief": self.cache["belief_content"][idx], "valid": self.cache["valid_frames"][idx], "danger_pf": self.cache["danger_pf"][idx], "hazard": self.hazard[idx], "tick_action": int(self.cache["tick_action"][idx]), } def collate(batch): return { "belief": torch.stack([b["belief"] for b in batch]), "valid": torch.stack([b["valid"] for b in batch]), "danger_pf": torch.stack([b["danger_pf"] for b in batch]), "hazard": torch.stack([b["hazard"] for b in batch]), "tick_action": torch.tensor([b["tick_action"] for b in batch], dtype=torch.long), } @torch.no_grad() def evaluate(model, loader, device): model.eval() all_hazard_logits, all_hazard_t, all_alert_score, all_alert_t = [], [], [], [] all_pf_logit, all_danger_pf, all_valid = [], [], [] for b in loader: bc = b["belief"].to(device, dtype=torch.float32) v = b["valid"].to(device) out = model(bc, valid_frames=v) all_hazard_logits.append(out["hazard_logits"].cpu().numpy()) all_hazard_t.append(b["hazard"].numpy()) all_alert_score.append(out["clip"].cpu().numpy()) # alert ground-truth = (tick_action == 2) all_alert_t.append((b["tick_action"] == 2).numpy().astype(int)) all_pf_logit.append(out["per_frame_logits"].cpu().numpy()) all_danger_pf.append(b["danger_pf"].numpy()) all_valid.append(v.cpu().numpy()) hz_logits = np.concatenate(all_hazard_logits) hz_t = np.concatenate(all_hazard_t) hz_pred = hz_logits.argmax(axis=-1) a_s = np.concatenate(all_alert_score) a_t = np.concatenate(all_alert_t) metrics = { "hazard_acc": float(accuracy_score(hz_t, hz_pred)), "hazard_balanced_acc": float(balanced_accuracy_score(hz_t, hz_pred)), "alert_AP": float(average_precision_score(a_t, a_s)), "alert_AUROC": float(roc_auc_score(a_t, a_s)), } return metrics def main(): ap = argparse.ArgumentParser(description=__doc__) ap.add_argument("--train_cache", type=Path, default=ROOT / "data/belief_cache_v3/sft_x_v3__train_9k.pt") ap.add_argument("--val_cache", type=Path, default=ROOT / "data/belief_cache_v3/sft_x_v3__multisrc_val.pt") ap.add_argument("--train_hazard", type=Path, default=ROOT / "data/policy_labels/hazard_categories_train_9k.json") ap.add_argument("--val_hazard", type=Path, default=ROOT / "data/policy_labels/hazard_categories_multisrc_val.json") ap.add_argument("--out_dir", type=Path, default=ROOT / "checkpoints/danger_v3_hazard") ap.add_argument("--in_dim", type=int, default=10240) ap.add_argument("--hidden", type=int, default=512) ap.add_argument("--k_queries", type=int, default=4) ap.add_argument("--dropout", type=float, default=0.2) ap.add_argument("--lr", type=float, default=5e-4) ap.add_argument("--weight_decay", type=float, default=1e-4) ap.add_argument("--epochs", type=int, default=50) ap.add_argument("--batch_size", type=int, default=64) ap.add_argument("--hazard_weight", type=float, default=0.3) ap.add_argument("--w_clip", type=float, default=0.5) ap.add_argument("--patience", type=int, default=15) ap.add_argument("--seed", type=int, default=0) ap.add_argument("--max_samples", type=int, default=0, help="if >0, truncate train+val for smoke testing") args = ap.parse_args() args.out_dir.mkdir(parents=True, exist_ok=True) torch.manual_seed(args.seed) device = "cuda" if torch.cuda.is_available() else "cpu" train_ds = HazardDataset(args.train_cache, args.train_hazard) val_ds = HazardDataset(args.val_cache, args.val_hazard) if args.max_samples > 0: train_ds.cache["ids"] = train_ds.cache["ids"][:args.max_samples] train_ds.hazard = train_ds.hazard[:args.max_samples] val_ds.cache["ids"] = val_ds.cache["ids"][:args.max_samples] val_ds.hazard = val_ds.hazard[:args.max_samples] train_loader = DataLoader(train_ds, batch_size=args.batch_size, shuffle=True, num_workers=2, collate_fn=collate, pin_memory=True) val_loader = DataLoader(val_ds, batch_size=args.batch_size * 2, shuffle=False, num_workers=2, collate_fn=collate, pin_memory=True) model = DangerHead(in_dim=args.in_dim, hidden=args.hidden, k_queries=args.k_queries, dropout=args.dropout, n_hazards=N_HAZARDS).to(device) logger.info(f" DangerHead-v3-hazard: " f"{sum(p.numel() for p in model.parameters())/1e6:.2f}M params") opt = torch.optim.AdamW(model.parameters(), lr=args.lr, weight_decay=args.weight_decay) n_steps = args.epochs * len(train_loader) sched = torch.optim.lr_scheduler.CosineAnnealingLR(opt, T_max=n_steps) log_records = [] best_score = -1e9 bad_epochs = 0 for ep in range(args.epochs): model.train() run = {"loss": 0, "danger": 0, "hazard": 0}; n_b = 0 pbar = tqdm(train_loader, ncols=80, desc=f"ep{ep}") for b in pbar: bc = b["belief"].to(device, dtype=torch.float32, non_blocking=True) v = b["valid"].to(device, non_blocking=True) dpf = b["danger_pf"].to(device, non_blocking=True) hz = b["hazard"].to(device, non_blocking=True) out = model(bc, valid_frames=v) d_l = danger_loss(out, dpf, valid_frames=v, w_clip=args.w_clip) h_l = F.cross_entropy(out["hazard_logits"], hz) total = d_l["loss"] + args.hazard_weight * h_l total.backward() torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) opt.step(); sched.step(); opt.zero_grad(set_to_none=True) run["loss"] += total.item() run["danger"] += d_l["loss"].item() run["hazard"] += h_l.item() n_b += 1 pbar.set_postfix(loss=run["loss"]/n_b, hz=run["hazard"]/n_b) m = evaluate(model, val_loader, device) rec = {"ep": ep, "train_loss": run["loss"]/max(1, n_b), "train_danger": run["danger"]/max(1, n_b), "train_hazard": run["hazard"]/max(1, n_b), "val": m} log_records.append(rec) logger.info(f"[ep{ep}] train={rec['train_loss']:.4f} " f"haz={rec['train_hazard']:.4f} | " f"val: alert_AP={m['alert_AP']:.4f} " f"alert_AUROC={m['alert_AUROC']:.4f} " f"hazard_bal_acc={m['hazard_balanced_acc']:.4f}") # Composite: 0.5 alert_AP + 0.3 alert_AUROC + 0.2 hazard_bal_acc score = (0.5 * m["alert_AP"] + 0.3 * m["alert_AUROC"] + 0.2 * m["hazard_balanced_acc"]) if score > best_score: best_score = score; bad_epochs = 0 save_dict = { "model": model.state_dict(), "in_dim": args.in_dim, "hidden": args.hidden, "k_queries": args.k_queries, "dropout": args.dropout, "n_hazards": N_HAZARDS, "val_metrics": m, "composite": score, "epoch": ep, "args": vars(args), } torch.save(save_dict, args.out_dir / "best.pt") logger.info(f" [save best] composite={score:.4f}") else: bad_epochs += 1 if bad_epochs >= args.patience: logger.info(f" early stop @ ep{ep} (patience {args.patience})") break (args.out_dir / "training_log.json").write_text( json.dumps(log_records, indent=2, default=str)) logger.info(f"\n[done] best composite={best_score:.4f} saved to {args.out_dir}/best.pt") if __name__ == "__main__": main()