#!/usr/bin/env python3 """ Fine-tune MViT-v2-s (Multiscale Vision Transformer) on Nexar collision data. Architecture: torchvision.models.video.mvit_v2_s (pretrained Kinetics-400) - Replace head (head.proj) with Linear(768, 1) for binary classification - Full fine-tuning with low LR for backbone, higher LR for head This replicates the 1st-place winning approach (0.898 mAP on private LB). Usage: python -m training.Nexar.mvit_trainer \ --train_csv nexar-collision-prediction/train.csv \ --video_dir nexar-collision-prediction/train \ --output_dir checkpoints/Nexar/mvit_v1 \ --epochs 20 \ --batch_size 8 \ --min_warning 0.3 # Data-centric ablation (more aggressive filtering): python -m training.Nexar.mvit_trainer \ --train_csv nexar-collision-prediction/train.csv \ --video_dir nexar-collision-prediction/train \ --output_dir checkpoints/Nexar/mvit_v2_strict \ --min_warning 1.0 \ --epochs 25 """ from __future__ import annotations import argparse import json import logging import random from pathlib import Path from typing import List import numpy as np import pandas as pd 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 tqdm import tqdm import sys sys.path.insert(0, str(Path(__file__).resolve().parent.parent.parent)) from training.Nexar.mvit_dataset import NexarMViTDataset, make_train_val_split logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s") logger = logging.getLogger("Nexar.mvit_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 build_mvit(pretrained: bool = True) -> nn.Module: """Load MViT-v2-s and replace head for binary classification.""" from torchvision.models.video import mvit_v2_s, MViT_V2_S_Weights weights = MViT_V2_S_Weights.DEFAULT if pretrained else None model = mvit_v2_s(weights=weights) # Replace classification head (Linear(768, 400) → Linear(768, 1)) in_features = model.head[1].in_features model.head[1] = nn.Linear(in_features, 1) nn.init.normal_(model.head[1].weight, std=0.01) nn.init.zeros_(model.head[1].bias) total = sum(p.numel() for p in model.parameters()) logger.info(f"MViT-v2-s total params: {total/1e6:.1f}M head_features: {in_features}") return model def make_sampler(labels: List[int]) -> WeightedRandomSampler: 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 * max(n_pos, 1)), len(labels_arr) / (2 * max(n_neg, 1))) return WeightedRandomSampler( weights=torch.from_numpy(weights).float(), num_samples=len(labels), replacement=True, ) def train_epoch(model, loader, optimizer, scaler, device) -> float: model.train() total_loss = 0.0 n = 0 for batch in tqdm(loader, desc="Train", leave=False): videos = batch["video"].to(device) # [B, C, T, H, W] labels = batch["label"].to(device) # [B] with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): logits = model(videos).squeeze(-1) # [B] loss = F.binary_cross_entropy_with_logits(logits, labels) scaler.scale(loss).backward() scaler.unscale_(optimizer) torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0) scaler.step(optimizer) scaler.update() optimizer.zero_grad() total_loss += loss.item() * len(labels) n += len(labels) return total_loss / max(n, 1) @torch.no_grad() def eval_epoch(model, loader, device): from sklearn.metrics import average_precision_score, roc_auc_score model.eval() all_scores: List[float] = [] all_labels: List[float] = [] total_loss = 0.0 n = 0 for batch in tqdm(loader, desc="Val", leave=False): videos = batch["video"].to(device) labels = batch["label"].to(device) with torch.amp.autocast(device_type="cuda", dtype=torch.bfloat16): logits = model(videos).squeeze(-1) loss = F.binary_cross_entropy_with_logits(logits, labels) scores = torch.sigmoid(logits) 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) try: ap = float(average_precision_score(arr_l, arr_s)) auc = float(roc_auc_score(arr_l, arr_s)) except Exception: ap = auc = float("nan") return total_loss / max(n, 1), ap, auc def main(): parser = argparse.ArgumentParser("mvit_trainer") parser.add_argument("--train_csv", default="nexar-collision-prediction/train.csv") parser.add_argument("--video_dir", default="nexar-collision-prediction/train", help="Root dir with {vid_id}.mp4 train videos") parser.add_argument("--output_dir", required=True) parser.add_argument("--pos_subdir", default="", help="If positive videos are in a subdirectory (e.g. 'positive')") parser.add_argument("--neg_subdir", default="", help="If negative videos are in a subdirectory (e.g. 'negative')") parser.add_argument("--epochs", type=int, default=20) parser.add_argument("--batch_size", type=int, default=8) parser.add_argument("--lr", type=float, default=5e-5, help="LR for backbone; head LR = lr * 10") parser.add_argument("--lr_min", type=float, default=1e-7) parser.add_argument("--weight_decay",type=float, default=1e-4) parser.add_argument("--val_frac", type=float, default=0.15) parser.add_argument("--min_warning", type=float, default=0.3, help="Data-centric filter: skip positives with warning < this (seconds)") parser.add_argument("--patience", type=int, default=6) parser.add_argument("--n_frames", type=int, default=16) parser.add_argument("--img_size", type=int, default=224) parser.add_argument("--no_pretrain", action="store_true") 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 split ──────────────────────────────────────────────────────────── train_df, val_df = make_train_val_split( args.train_csv, args.val_frac, min_warning_s=args.min_warning, ) train_csv_path = out_dir / "_train_split.csv" val_csv_path = out_dir / "_val_split.csv" train_df.to_csv(train_csv_path, index=False) val_df.to_csv(val_csv_path, index=False) train_ds = NexarMViTDataset( str(train_csv_path), args.video_dir, train_mode=True, pos_subdir=args.pos_subdir, neg_subdir=args.neg_subdir, min_warning_s=args.min_warning, n_frames=args.n_frames, img_size=args.img_size, ) val_ds = NexarMViTDataset( str(val_csv_path), args.video_dir, train_mode=False, pos_subdir=args.pos_subdir, neg_subdir=args.neg_subdir, min_warning_s=0.0, # no filter on validation n_frames=args.n_frames, img_size=args.img_size, ) train_labels = [s["label"] for s in train_ds.samples] sampler = make_sampler(train_labels) train_loader = DataLoader( train_ds, batch_size=args.batch_size, sampler=sampler, num_workers=4, pin_memory=True, drop_last=True, ) val_loader = DataLoader( val_ds, batch_size=args.batch_size, shuffle=False, num_workers=4, pin_memory=True, ) # ── model ───────────────────────────────────────────────────────────────── model = build_mvit(pretrained=not args.no_pretrain).to(device) # Differential learning rates: higher LR for head head_params = list(model.head.parameters()) head_ids = {id(p) for p in head_params} backbone_params = [p for p in model.parameters() if id(p) not in head_ids] optimizer = AdamW([ {"params": backbone_params, "lr": args.lr}, {"params": head_params, "lr": args.lr * 10}, ], weight_decay=args.weight_decay) total_steps = args.epochs * len(train_loader) scheduler = CosineAnnealingLR(optimizer, T_max=total_steps, eta_min=args.lr_min) scaler = torch.amp.GradScaler() # ── 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, scaler, device) scheduler.step() val_loss, val_ap, val_auc = eval_epoch(model, val_loader, device) lr_bb = 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_bb:.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, "n_frames": args.n_frames, "img_size": args.img_size, "min_warning": args.min_warning, "model": "mvit_v2_s", }, f, indent=2) logger.info(f" ★ New best val_AP={best_ap:.4f}") else: patience_count += 1 if patience_count >= args.patience: logger.info(f"Early stopping at epoch {epoch}") break with open(out_dir / "history.json", "w") as f: json.dump(history, f, indent=2) logger.info(f"\n✅ Done. Best val_AP = {best_ap:.4f}") logger.info(f" Checkpoint: {out_dir}/best_model.pt") if __name__ == "__main__": main()