icassp2027-experiment/train_turn_head.py

"""Train the frozen-Omni EasyCom anticipation head.

The defaults follow EgoSpeak's EasyCom LSTR optimization config where the
architectures are comparable: 50 epochs, Adam, 7e-5 learning rate, 5e-5 weight
decay, linear warmup for 10 epochs, cosine decay, dropout 0.2, and MCE loss.
"""

from __future__ import annotations

import argparse
import json
import math
import random
from pathlib import Path
from typing import Any

import numpy as np
import torch
from torch.optim.lr_scheduler import LambdaLR
from torch.utils.data import DataLoader, TensorDataset
from tqdm.auto import tqdm

from easycom_head import (
    CLASS_NAMES,
    NUM_CLASSES,
    NUM_HORIZONS,
    OFFICIAL_TRAIN_SESSIONS,
    STEP_SECONDS,
    TurnPredictionHead,
    average_precision_metrics,
    egospeak_multihorizon_loss,
    load_session_data,
    miniroad_multihorizon_loss,
    normalize_session_name,
    validate_disjoint_sessions,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--feature-root", type=Path, required=True)
    parser.add_argument("--target-root", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    parser.add_argument(
        "--sessions",
        nargs="+",
        default=list(OFFICIAL_TRAIN_SESSIONS),
        help="Training sessions; default is the official nine-session split.",
    )
    parser.add_argument(
        "--dev-sessions",
        nargs="*",
        default=[],
        help=(
            "Optional held-out training sessions used for checkpoint selection. "
            "They are automatically removed from --sessions. Never use 1/2/3."
        ),
    )
    parser.add_argument("--epochs", type=int, default=50)
    parser.add_argument("--batch-size", type=int, default=512)
    parser.add_argument("--learning-rate", type=float, default=7e-5)
    parser.add_argument("--weight-decay", type=float, default=5e-5)
    parser.add_argument("--warmup-epochs", type=float, default=10.0)
    parser.add_argument("--warmup-factor", type=float, default=0.3)
    parser.add_argument("--dropout", type=float, default=0.2)
    parser.add_argument(
        "--loss",
        choices=("lstr-mce", "miniroad-normalized"),
        default="lstr-mce",
    )
    parser.add_argument("--device", default="cuda:0")
    parser.add_argument("--num-workers", type=int, default=0)
    parser.add_argument("--seed", type=int, default=0)
    parser.add_argument(
        "--expected-mode",
        default="causal-prefix",
        help="Expected feature mode, or 'any' to disable the check.",
    )
    parser.add_argument("--no-progress", action="store_true")
    return parser.parse_args()


def _validate_args(args: argparse.Namespace) -> None:
    if args.epochs <= 0:
        raise ValueError("--epochs must be positive")
    if args.batch_size <= 0:
        raise ValueError("--batch-size must be positive")
    if args.learning_rate <= 0:
        raise ValueError("--learning-rate must be positive")
    if args.weight_decay < 0:
        raise ValueError("--weight-decay cannot be negative")
    if not 0 <= args.warmup_epochs <= args.epochs:
        raise ValueError("--warmup-epochs must be in [0, epochs]")
    if not 0 < args.warmup_factor <= 1:
        raise ValueError("--warmup-factor must be in (0, 1]")
    if not 0 <= args.dropout < 1:
        raise ValueError("--dropout must be in [0, 1)")
    if args.num_workers < 0:
        raise ValueError("--num-workers cannot be negative")


def _resolve_device(name: str) -> torch.device:
    device = torch.device(name)
    if device.type == "cuda" and not torch.cuda.is_available():
        raise RuntimeError(f"CUDA requested but unavailable: {name}")
    return device


def _seed_everything(seed: int) -> None:
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed_all(seed)


def _load_examples(
    feature_root: Path,
    target_root: Path,
    sessions: list[str],
    expected_mode: str | None,
) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
    all_states: list[torch.Tensor] = []
    all_targets: list[torch.Tensor] = []
    metadata: list[dict[str, Any]] = []
    hidden_size: int | None = None
    for name in sessions:
        data = load_session_data(
            feature_root,
            target_root,
            name,
            expected_mode=expected_mode,
        )
        states, targets, _ = data.future_examples(NUM_HORIZONS)
        if hidden_size is None:
            hidden_size = data.hidden_size
        elif data.hidden_size != hidden_size:
            raise ValueError(
                f"Feature dimensions differ: {name} has {data.hidden_size}, "
                f"expected {hidden_size}"
            )
        all_states.append(states)
        all_targets.append(targets)
        metadata.append(
            {
                "session": name,
                "num_states": len(data.states),
                "num_training_examples": len(states),
                **data.feature_metadata,
            }
        )
        print(
            f"{name}: states={len(data.states)}, "
            f"valid_10_horizon_examples={len(states)}"
        )
    if hidden_size is None:
        raise ValueError("No sessions selected")
    return (
        torch.cat(all_states, dim=0),
        torch.cat(all_targets, dim=0),
        hidden_size,
        metadata,
    )


def _scheduler(
    optimizer: torch.optim.Optimizer,
    *,
    steps_per_epoch: int,
    epochs: int,
    warmup_epochs: float,
    warmup_factor: float,
) -> LambdaLR:
    total_steps = max(1, steps_per_epoch * epochs)
    warmup_steps = steps_per_epoch * warmup_epochs

    def multiplier(step: int) -> float:
        cosine = 0.5 * (1.0 + math.cos(math.pi * step / total_steps))
        if warmup_steps > 0 and step < warmup_steps:
            alpha = step / warmup_steps
            warmup = warmup_factor * (1.0 - alpha) + alpha
        else:
            warmup = 1.0
        return warmup * cosine

    return LambdaLR(optimizer, multiplier)


@torch.inference_mode()
def _evaluate_valid_horizons(
    head: TurnPredictionHead,
    states: torch.Tensor,
    targets: torch.Tensor,
    device: torch.device,
    batch_size: int,
) -> dict[str, Any]:
    head.eval()
    chunks: list[torch.Tensor] = []
    for start in range(0, len(states), batch_size):
        x = states[start : start + batch_size].to(
            device=device,
            dtype=torch.float32,
            non_blocking=True,
        )
        chunks.append(head(x).softmax(dim=-1).cpu())
    probabilities = torch.cat(chunks, dim=0).numpy()
    target_array = targets.numpy()

    results = []
    for horizon in range(NUM_HORIZONS):
        result = average_precision_metrics(
            probabilities[:, horizon],
            target_array[:, horizon],
        )
        result["horizon"] = horizon + 1
        result["seconds"] = (horizon + 1) * STEP_SECONDS
        results.append(result)
    return {
        "horizons": results,
        "mean_map": float(np.mean([item["map"] for item in results])),
    }


def _checkpoint(
    head: TurnPredictionHead,
    *,
    epoch: int,
    best_epoch: int,
    best_dev_map: float | None,
    train_loss: float,
    args: argparse.Namespace,
    train_sessions: list[str],
    dev_sessions: list[str],
    hidden_size: int,
    feature_metadata: list[dict[str, Any]],
    history: list[dict[str, Any]],
) -> dict[str, Any]:
    return {
        "format_version": 1,
        "head": head.state_dict(),
        "head_config": {
            "hidden_size": hidden_size,
            "num_horizons": NUM_HORIZONS,
            "num_classes": NUM_CLASSES,
            "dropout": args.dropout,
        },
        "alignment": {
            "state_label_offset": 0,
            "future_offsets": list(range(1, NUM_HORIZONS + 1)),
            "future_seconds": [
                (index + 1) * STEP_SECONDS for index in range(NUM_HORIZONS)
            ],
        },
        "classes": list(CLASS_NAMES),
        "epoch": epoch,
        "best_epoch": best_epoch,
        "best_dev_map": best_dev_map,
        "train_loss": train_loss,
        "train_sessions": train_sessions,
        "dev_sessions": dev_sessions,
        "feature_metadata": feature_metadata,
        "training_config": {
            "epochs": args.epochs,
            "batch_size": args.batch_size,
            "learning_rate": args.learning_rate,
            "weight_decay": args.weight_decay,
            "warmup_epochs": args.warmup_epochs,
            "warmup_factor": args.warmup_factor,
            "loss": args.loss,
            "optimizer": "Adam",
            "scheduler": "linear_warmup_times_cosine",
            "seed": args.seed,
            "expected_mode": args.expected_mode,
        },
        "history": history,
    }


def main() -> None:
    args = parse_args()
    _validate_args(args)
    _seed_everything(args.seed)
    device = _resolve_device(args.device)
    expected_mode = None if args.expected_mode == "any" else args.expected_mode

    requested_train = [normalize_session_name(item) for item in args.sessions]
    dev_sessions = [normalize_session_name(item) for item in args.dev_sessions]
    train_sessions = [item for item in requested_train if item not in dev_sessions]
    validate_disjoint_sessions(train_sessions, dev_sessions)
    if not train_sessions:
        raise ValueError("No training sessions remain after removing dev sessions")
    forbidden_test = {"Session_1", "Session_2", "Session_3"}
    leaked = forbidden_test & (set(train_sessions) | set(dev_sessions))
    if leaked:
        raise ValueError(
            f"Official test sessions cannot be used for training/selection: "
            f"{sorted(leaked)}"
        )

    print(f"device: {device}")
    print(f"train sessions: {', '.join(train_sessions)}")
    print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}")
    states, targets, hidden_size, feature_metadata = _load_examples(
        args.feature_root,
        args.target_root,
        train_sessions,
        expected_mode,
    )
    print(f"training examples: {len(states)}, feature_dim={hidden_size}")
    dev_states: torch.Tensor | None = None
    dev_targets: torch.Tensor | None = None
    if dev_sessions:
        dev_states, dev_targets, dev_hidden_size, _ = _load_examples(
            args.feature_root,
            args.target_root,
            dev_sessions,
            expected_mode,
        )
        if dev_hidden_size != hidden_size:
            raise ValueError(
                f"Dev feature dimension {dev_hidden_size} differs from "
                f"training dimension {hidden_size}"
            )
        print(f"development examples: {len(dev_states)}")

    dataset = TensorDataset(states, targets)
    generator = torch.Generator().manual_seed(args.seed)
    loader = DataLoader(
        dataset,
        batch_size=args.batch_size,
        shuffle=True,
        num_workers=args.num_workers,
        pin_memory=device.type == "cuda",
        persistent_workers=args.num_workers > 0,
        generator=generator,
    )
    head = TurnPredictionHead(
        hidden_size=hidden_size,
        num_horizons=NUM_HORIZONS,
        num_classes=NUM_CLASSES,
        dropout=args.dropout,
    ).to(device)
    optimizer = torch.optim.Adam(
        head.parameters(),
        lr=args.learning_rate,
        weight_decay=args.weight_decay,
    )
    scheduler = _scheduler(
        optimizer,
        steps_per_epoch=len(loader),
        epochs=args.epochs,
        warmup_epochs=args.warmup_epochs,
        warmup_factor=args.warmup_factor,
    )
    loss_function = (
        egospeak_multihorizon_loss
        if args.loss == "lstr-mce"
        else miniroad_multihorizon_loss
    )

    args.output.parent.mkdir(parents=True, exist_ok=True)
    history: list[dict[str, Any]] = []
    best_dev_map = -math.inf if dev_sessions else None
    best_epoch = args.epochs if not dev_sessions else 0
    for epoch in range(1, args.epochs + 1):
        head.train()
        loss_sum = 0.0
        example_count = 0
        progress = tqdm(
            loader,
            desc=f"head epoch {epoch}/{args.epochs}",
            unit="batch",
            disable=args.no_progress,
        )
        for state_batch, target_batch in progress:
            state_batch = state_batch.to(
                device=device,
                dtype=torch.float32,
                non_blocking=True,
            )
            target_batch = target_batch.to(
                device=device,
                dtype=torch.float32,
                non_blocking=True,
            )
            optimizer.zero_grad(set_to_none=True)
            logits = head(state_batch)
            loss = loss_function(logits, target_batch)
            loss.backward()
            optimizer.step()
            scheduler.step()
            batch_count = len(state_batch)
            loss_sum += float(loss.detach()) * batch_count
            example_count += batch_count
            progress.set_postfix(
                loss=f"{float(loss.detach()):.5f}",
                lr=f"{scheduler.get_last_lr()[0]:.3g}",
            )
        train_loss = loss_sum / example_count
        record: dict[str, Any] = {
            "epoch": epoch,
            "train_loss": train_loss,
            "learning_rate": scheduler.get_last_lr()[0],
        }

        if dev_sessions:
            assert dev_states is not None and dev_targets is not None
            dev_metrics = _evaluate_valid_horizons(
                head,
                dev_states,
                dev_targets,
                device,
                args.batch_size,
            )
            dev_map = float(dev_metrics["mean_map"])
            record["dev_valid_mean_map"] = dev_map
            print(
                f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
                f"dev_valid_mean_mAP={dev_map * 100:.3f}"
            )
            if dev_map > float(best_dev_map):
                best_dev_map = dev_map
                best_epoch = epoch
                torch.save(
                    _checkpoint(
                        head,
                        epoch=epoch,
                        best_epoch=best_epoch,
                        best_dev_map=best_dev_map,
                        train_loss=train_loss,
                        args=args,
                        train_sessions=train_sessions,
                        dev_sessions=dev_sessions,
                        hidden_size=hidden_size,
                        feature_metadata=feature_metadata,
                        history=history + [record],
                    ),
                    args.output,
                )
                print(f"saved new best checkpoint: {args.output}")
        else:
            print(f"epoch {epoch:02d}: train_loss={train_loss:.6f}")
        history.append(record)

    if not dev_sessions:
        torch.save(
            _checkpoint(
                head,
                epoch=args.epochs,
                best_epoch=best_epoch,
                best_dev_map=None,
                train_loss=history[-1]["train_loss"],
                args=args,
                train_sessions=train_sessions,
                dev_sessions=dev_sessions,
                hidden_size=hidden_size,
                feature_metadata=feature_metadata,
                history=history,
            ),
            args.output,
        )
    print(f"checkpoint: {args.output}")
    if dev_sessions:
        print(
            f"best epoch: {best_epoch}, "
            f"dev valid-only mean mAP: {float(best_dev_map) * 100:.3f}"
        )
    summary_path = args.output.with_suffix(args.output.suffix + ".json")
    payload = torch.load(args.output, map_location="cpu", weights_only=False)
    summary = {key: value for key, value in payload.items() if key != "head"}
    summary_path.write_text(json.dumps(summary, indent=2) + "\n")
    print(f"training summary: {summary_path}")


if __name__ == "__main__":
    main()