"""Train the frozen-Omni EasyCom or Ego4D 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, DATASET_NAMES, NUM_CLASSES, NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS, STEP_SECONDS, TurnPredictionHead, average_precision_metrics, egospeak_multihorizon_loss, load_ego4d_manifest_split, load_session_data, miniroad_multihorizon_loss, normalize_sample_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( "--dataset", choices=DATASET_NAMES, default="easycom", help="Dataset naming/split contract (default: easycom).", ) parser.add_argument( "--manifest", type=Path, help="Required for Ego4D; TSV containing split and uid columns.", ) parser.add_argument( "--sessions", nargs="+", default=None, help=( "Training samples. Defaults to the official EasyCom sessions or " "all Ego4D manifest train UIDs." ), ) 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( "--sparse-anchor-cache-dir", type=Path, action="append", default=[], help=( "Optional transition-sparse AV cache directory containing " "manifest.json. Repeat once per training session. Only its anchor " "target_index values select examples; states still come from " "--feature-root. Development examples always remain dense." ), ) 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") if args.dataset == "ego4d" and args.manifest is None: raise ValueError("--manifest is required when --dataset ego4d") 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, dataset: str, selected_indices: dict[str, set[int]] | None = 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, dataset=dataset, ) states, targets, future_indices = data.future_examples(NUM_HORIZONS) dense_example_count = len(states) if selected_indices is not None: if name not in selected_indices: raise ValueError(f"No sparse anchor manifest was supplied for {name}") wanted = selected_indices[name] current_indices = future_indices[:, 0] - 1 keep = torch.tensor( [int(index) in wanted for index in current_indices.tolist()], dtype=torch.bool, ) states = states[keep] targets = targets[keep] if not len(states): raise ValueError(f"Sparse anchor selection kept no examples for {name}") 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), "num_dense_training_examples": dense_example_count, "supervision_policy": ( "transition-sparse-manifest" if selected_indices is not None else "dense-5hz" ), **data.feature_metadata, } ) print( f"{name}: states={len(data.states)}, " f"valid_10_horizon_examples={dense_example_count}, " f"selected_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 _load_sparse_anchor_indices( cache_dirs: list[Path], ) -> dict[str, set[int]] | None: if not cache_dirs: return None result: dict[str, set[int]] = {} for cache_dir in cache_dirs: manifest_path = cache_dir.expanduser().resolve() / "manifest.json" if not manifest_path.is_file(): raise FileNotFoundError(manifest_path) manifest = json.loads(manifest_path.read_text()) session = normalize_sample_name(str(manifest.get("session", "")), "easycom") policy = manifest.get("anchor_selection", {}).get("policy") if policy != "transition-sparse": raise ValueError( f"Expected transition-sparse manifest, got {policy!r}: " f"{manifest_path}" ) if session in result: raise ValueError(f"Duplicate sparse anchor manifest for {session}") indices = { int(anchor["target_index"]) for anchor in manifest.get("anchors", []) if bool(anchor.get("has_future_target", True)) } if not indices: raise ValueError(f"No valid sparse anchors in {manifest_path}") result[session] = indices print(f"sparse anchors {session}: {len(indices)} from {manifest_path}") return result 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), "dataset": args.dataset, "manifest": ( str(args.manifest.expanduser().resolve()) if args.manifest is not None else None ), "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, "supervision_policy": ( "transition-sparse-manifest" if args.sparse_anchor_cache_dir else "dense-5hz" ), "sparse_anchor_cache_dirs": [ str(path.expanduser().resolve()) for path in args.sparse_anchor_cache_dir ], }, "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 sparse_anchor_indices = _load_sparse_anchor_indices( args.sparse_anchor_cache_dir ) if sparse_anchor_indices is not None and args.dataset != "easycom": raise ValueError("--sparse-anchor-cache-dir currently supports EasyCom only") if args.dataset == "ego4d": assert args.manifest is not None official_train = load_ego4d_manifest_split(args.manifest, "train") official_test = set(load_ego4d_manifest_split(args.manifest, "test")) requested_values = args.sessions or official_train else: official_train = list(OFFICIAL_TRAIN_SESSIONS) official_test = {"Session_1", "Session_2", "Session_3"} requested_values = args.sessions or official_train requested_train = [ normalize_sample_name(item, args.dataset) for item in requested_values ] dev_sessions = [ normalize_sample_name(item, args.dataset) for item in args.dev_sessions ] if len(requested_train) != len(set(requested_train)): raise ValueError("Duplicate training samples were supplied") if len(dev_sessions) != len(set(dev_sessions)): raise ValueError("Duplicate development samples were supplied") if args.dataset == "ego4d": allowed_train = set(official_train) unknown = (set(requested_train) | set(dev_sessions)) - allowed_train if unknown: raise ValueError( "Ego4D training/dev samples must belong to the manifest train " f"split; invalid: {sorted(unknown)}" ) 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") if sparse_anchor_indices is not None: missing = set(train_sessions) - set(sparse_anchor_indices) extra = set(sparse_anchor_indices) - set(train_sessions) if missing or extra: raise ValueError( "Sparse anchor manifests must match training sessions exactly; " f"missing={sorted(missing)}, extra={sorted(extra)}" ) leaked = official_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, args.dataset, sparse_anchor_indices, ) 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, args.dataset, None, ) 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()