training/Nexar/mvit_trainer.py

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