"""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 hashlib 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, HIGHRATE_PERSON_AV_FUSION_MODES, LAYER_FUSION_MODES, LOCAL_FUSION_MODES, LOCAL_FUSION_SCOPES, LOSS_NAMES, NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS, STEP_SECONDS, TRACK_PERSON_FUSION_MODES, TRANSITION_OBJECTIVE_NAMES, TurnPredictionHead, average_precision_metrics, checkpoint_head_config, checkpoint_loss_name, egospeak_multihorizon_loss, four_state_multihorizon_loss, head_output_names_for_loss, independent_bce_multihorizon_loss, load_ego4d_manifest_split, load_session_data, marginal_probabilities_from_logits, miniroad_multihorizon_loss, normalize_sample_name, output_classes_for_loss, output_space_for_loss, transition_aware_objective, transition_targets_from_activity, validate_disjoint_sessions, ) from local_turn_features import FramePersonFeatures INITIALIZED_PERSON_RESIDUAL_MODES = ( *TRACK_PERSON_FUSION_MODES, *HIGHRATE_PERSON_AV_FUSION_MODES, ) 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( "--local-feature-root", type=Path, help="Optional frozen local-feature sidecars; omit for the original baseline.", ) parser.add_argument("--local-modalities", choices=("visual", "audio", "av"), default="av") parser.add_argument( "--local-fusion", choices=LOCAL_FUSION_MODES, default="pooled", help=( "pooled keeps the existing baseline; frame-set-attention is F1; " "track-gru-residual preserves the legacy zero-gated screen; " "causal-person-adapter-v2 adds a dynamic-gated, zero-output " "causal track correction to a matched pooled AV checkpoint; " "causal-person-adapter-egospeak-gru keeps that adapter but uses " "EgoSpeak MiniROADA's Linear-LN-ReLU-Dropout-GRU-ReLU temporal " "stem on each anonymous person track; " "highrate-person-av-attention lets global Omni plus one shared " "traditional CPC audio stream query causal 25 Hz person visual/ASD " "states sampled at the 5 Hz cutoffs; " "audio-person-frame-set and gated-person-state-adapter are " "current-sidecar F5/F6 screening variants, not full F5/F6." ), ) parser.add_argument( "--local-projection-dim", type=int, default=256, help="Trainable projection per local modality; 0 retains encoder dimensions.", ) parser.add_argument("--person-attention-dim", type=int, default=256) parser.add_argument("--person-attention-heads", type=int, default=4) parser.add_argument( "--person-history-frames", type=int, default=16, help="Maximum causal observations per current face track (16 at 2 FPS = 8s).", ) parser.add_argument( "--person-track-max-gap-seconds", type=float, default=1.0, help="Reset a provided track when consecutive observations are farther apart.", ) parser.add_argument( "--initialize-pooled-checkpoint", type=Path, help=( "Required by person residual modes: copy an exact pooled G+A(+V) " "base before training the person correction." ), ) parser.add_argument( "--train-pooled-base", action="store_true", help="Also update the initialized pooled base; default keeps it frozen.", ) 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( "--learnable-layer-mixture", action="store_true", help=( "Legacy alias for --layer-fusion scalar-mixture." ), ) parser.add_argument( "--layer-fusion", choices=tuple(mode for mode in LAYER_FUSION_MODES if mode != "single"), default=None, help=( "Fusion for stacked [T,L,D] query features. fixed-mean has no " "fusion parameters; scalar-mixture learns L global weights; " "concat-mlp flattens L*D and learns a full projection to D." ), ) parser.add_argument( "--loss", choices=LOSS_NAMES, default="lstr-mce", help=( "lstr-mce/miniroad-normalized use three competing softmax logits; " "independent-bce uses three sigmoid marginals; four-state-ce uses " "silence/wearer-only/others-only/overlap softmax logits." ), ) parser.add_argument( "--transition-objective", choices=TRANSITION_OBJECTIVE_NAMES, default="none", help=( "none keeps the activity-only baseline; r1-supervised adds a " "weighted binary transition head; r2-expected-reward additionally " "maximizes exact expected transition correctness." ), ) parser.add_argument("--transition-loss-weight", type=float, default=0.2) parser.add_argument("--transition-pos-weight", type=float, default=3.0) parser.add_argument("--transition-reward-weight", type=float, default=0.1) 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.local_projection_dim < 0: raise ValueError("--local-projection-dim must be non-negative") if args.person_attention_dim <= 0 or args.person_attention_heads <= 0: raise ValueError("Person attention dimension/heads must be positive") if args.person_attention_dim % args.person_attention_heads: raise ValueError("--person-attention-dim must be divisible by --person-attention-heads") if args.person_history_frames <= 0: raise ValueError("--person-history-frames must be positive") if args.person_track_max_gap_seconds <= 0: raise ValueError("--person-track-max-gap-seconds must be positive") if args.local_fusion != "pooled" and ( args.local_feature_root is None or args.local_modalities != "av" ): raise ValueError( "Person-set --local-fusion modes require --local-feature-root and " "--local-modalities av" ) if args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES: if args.initialize_pooled_checkpoint is None: raise ValueError( f"{args.local_fusion} requires --initialize-pooled-checkpoint" ) elif args.initialize_pooled_checkpoint is not None or args.train_pooled_base: raise ValueError( "--initialize-pooled-checkpoint/--train-pooled-base require " "a person-residual --local-fusion mode" ) 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.transition_loss_weight < 0: raise ValueError("--transition-loss-weight cannot be negative") if args.transition_pos_weight <= 0: raise ValueError("--transition-pos-weight must be positive") if args.transition_reward_weight < 0: raise ValueError("--transition-reward-weight cannot be negative") if args.dataset == "ego4d" and args.manifest is None: raise ValueError("--manifest is required when --dataset ego4d") if ( args.learnable_layer_mixture and args.layer_fusion not in {None, "scalar-mixture"} ): raise ValueError( "--learnable-layer-mixture is an alias for " "--layer-fusion scalar-mixture and cannot select another fusion" ) 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, local_feature_root: Path | None = None, local_modalities: str = "av", local_contract: dict[str, Any] | None = None, compute_transition_targets: bool = True, local_fusion: str = "pooled", person_track_max_gap_seconds: float = 1.0, ) -> tuple[ torch.Tensor, torch.Tensor, torch.Tensor, int, int, list[dict[str, Any]], torch.Tensor, FramePersonFeatures | None, ]: all_states: list[torch.Tensor] = [] all_targets: list[torch.Tensor] = [] all_transition_targets: list[torch.Tensor] = [] all_local: list[torch.Tensor] = [] all_people: list[FramePersonFeatures] = [] metadata: list[dict[str, Any]] = [] hidden_size: int | None = None num_input_layers: int | None = None for name in sessions: data = load_session_data( feature_root, target_root, name, expected_mode=expected_mode, dataset=dataset, local_feature_root=local_feature_root, local_modalities=local_modalities, local_contract=local_contract, local_fusion=local_fusion, person_track_max_gap_seconds=person_track_max_gap_seconds, ) local_contract = data.feature_metadata.get("local_contract") states, targets, future_indices = data.future_examples(NUM_HORIZONS) current_indices = future_indices[:, 0] - 1 local = ( data.local_features[current_indices] if data.local_features is not None else torch.empty(len(states), 0) ) people = ( data.local_person_features.select_rows(current_indices) if data.local_person_features is not None else None ) if compute_transition_targets: activity_sequence = torch.cat( (data.targets[current_indices, None, :], targets), dim=1, ) transition_targets = transition_targets_from_activity( activity_sequence ) else: transition_targets = torch.zeros( targets.shape[:-1], dtype=targets.dtype, ) 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] keep = torch.tensor( [int(index) in wanted for index in current_indices.tolist()], dtype=torch.bool, ) states = states[keep] targets = targets[keep] transition_targets = transition_targets[keep] local = local[keep] if people is not None: people = people.select_rows(torch.nonzero(keep).flatten()) 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}" ) if num_input_layers is None: num_input_layers = data.num_input_layers elif data.num_input_layers != num_input_layers: raise ValueError( f"Feature layer counts differ: {name} has " f"{data.num_input_layers}, expected {num_input_layers}" ) all_states.append(states) all_targets.append(targets) all_transition_targets.append(transition_targets) all_local.append(local) if people is not None: all_people.append(people) metadata.append( { "session": name, "num_states": len(data.states), "num_training_examples": len(states), "num_dense_training_examples": dense_example_count, "num_positive_transition_targets": int( transition_targets.sum().item() ), "transition_target_count": int(transition_targets.numel()), "transition_positive_rate": float(transition_targets.mean()), "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 or num_input_layers is None: raise ValueError("No sessions selected") if all_people and len(all_people) != len(all_states): raise ValueError("Person sets are missing for some sessions") return ( torch.cat(all_states, dim=0), torch.cat(all_targets, dim=0), torch.cat(all_transition_targets, dim=0), hidden_size, num_input_layers, metadata, torch.cat(all_local, dim=0), FramePersonFeatures.concatenate(all_people) if all_people else None, ) 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, loss_name: str, local: torch.Tensor | None = None, people: FramePersonFeatures | None = None, ) -> 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, ) side = ( local[start : start + batch_size].to(device=device, dtype=torch.float32) if local is not None else None ) person_batch = None person_mask = None if people is not None: rows = torch.arange(start, min(start + batch_size, len(states))) if head.local_fusion in TRACK_PERSON_FUSION_MODES: person_batch, person_mask = people.padded_history( rows, max_observations=head.person_history_frames, ) else: person_batch, person_mask = people.padded(rows) person_batch = person_batch.to(device=device, dtype=torch.float32) person_mask = person_mask.to(device=device) chunks.append( marginal_probabilities_from_logits( head(x, side, person_batch, person_mask), loss_name, ).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, num_input_layers: int, feature_metadata: list[dict[str, Any]], history: list[dict[str, Any]], ) -> dict[str, Any]: return { "format_version": 7, "head": head.state_dict(), "head_config": { "hidden_size": hidden_size, "num_horizons": NUM_HORIZONS, "num_classes": head.num_classes, "dropout": args.dropout, "num_input_layers": num_input_layers, "layer_fusion": args.layer_fusion, "learnable_layer_mixture": args.learnable_layer_mixture, "transition_auxiliary": head.transition_auxiliary, "local_visual_dim": head.local_visual_dim, "local_audio_dim": head.local_audio_dim, "local_projection_dim": head.local_projection_dim, "local_fusion": head.local_fusion, "person_attention_dim": head.person_attention_dim, "person_attention_heads": head.person_attention_heads, "person_history_frames": head.person_history_frames, "person_track_max_gap_seconds": head.person_track_max_gap_seconds, }, "local_contract": feature_metadata[0].get("local_contract"), "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), "head_output_names": list(head_output_names_for_loss(args.loss)), "output_space": output_space_for_loss(args.loss), "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, "transition_objective": args.transition_objective, "transition_loss_weight": args.transition_loss_weight, "transition_pos_weight": args.transition_pos_weight, "transition_reward_weight": args.transition_reward_weight, "optimizer": "Adam", "scheduler": "linear_warmup_times_cosine", "seed": args.seed, "expected_mode": args.expected_mode, "layer_fusion": args.layer_fusion, "local_feature_root": str(args.local_feature_root) if args.local_feature_root else None, "local_modalities": args.local_modalities if args.local_feature_root else None, "local_fusion": args.local_fusion, "local_fusion_scope": LOCAL_FUSION_SCOPES[args.local_fusion], "local_projection_dim": args.local_projection_dim, "person_attention_dim": args.person_attention_dim, "person_attention_heads": args.person_attention_heads, "person_history_frames": args.person_history_frames, "person_track_max_gap_seconds": args.person_track_max_gap_seconds, "initialize_pooled_checkpoint": ( str(args.initialize_pooled_checkpoint.expanduser().resolve()) if args.initialize_pooled_checkpoint is not None else None ), "initialize_pooled_checkpoint_sha256": getattr( args, "initialize_pooled_checkpoint_sha256", None, ), "initialize_pooled_checkpoint_epoch": getattr( args, "initialize_pooled_checkpoint_epoch", None, ), "train_pooled_base": args.train_pooled_base, "pooled_base_frozen": bool( args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES and not args.train_pooled_base ), "learnable_layer_mixture": args.layer_fusion == "scalar-mixture", "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 _initialize_person_residual_from_pooled( head: TurnPredictionHead, checkpoint: Path, *, args: argparse.Namespace, train_sessions: list[str], dev_sessions: list[str], local_contract: dict[str, Any] | None, ) -> None: """Load and optionally freeze the exact pooled base in a person model.""" path = checkpoint.expanduser().resolve() if not path.is_file(): raise FileNotFoundError(path) payload = torch.load(path, map_location="cpu", weights_only=False) if not isinstance(payload, dict) or not isinstance(payload.get("head"), dict): raise ValueError(f"Invalid pooled checkpoint: {path}") config = checkpoint_head_config(payload) if config["local_fusion"] != "pooled": raise ValueError( "Person residual initialization requires a pooled checkpoint, got " f"{config['local_fusion']!r}: {path}" ) if checkpoint_loss_name(payload) != args.loss: raise ValueError("Pooled checkpoint loss differs from the person experiment") if payload.get("dataset", args.dataset) != args.dataset: raise ValueError("Pooled checkpoint dataset differs from the person experiment") if payload.get("train_sessions") != train_sessions: raise ValueError("Pooled checkpoint train sessions differ from this run") if payload.get("dev_sessions", []) != dev_sessions: raise ValueError("Pooled checkpoint dev sessions differ from this run") if args.local_fusion in HIGHRATE_PERSON_AV_FUSION_MODES: expected_base_contract = ( {**local_contract, "modalities": "audio"} if local_contract is not None else None ) expected_local_visual_dim = 0 else: expected_base_contract = local_contract expected_local_visual_dim = head.local_visual_dim if payload.get("local_contract") != expected_base_contract: raise ValueError("Pooled checkpoint local feature contract differs from this run") expected = { "hidden_size": head.hidden_size, "num_horizons": head.num_horizons, "num_classes": head.num_classes, "num_input_layers": head.num_input_layers, "layer_fusion": head.layer_fusion, "local_visual_dim": expected_local_visual_dim, "local_audio_dim": head.local_audio_dim, "local_projection_dim": head.local_projection_dim, "dropout": head.dropout, "transition_auxiliary": head.transition_auxiliary, } mismatches = { key: (config.get(key), value) for key, value in expected.items() if config.get(key) != value } if mismatches: raise ValueError( f"Pooled checkpoint architecture differs from person base: {mismatches}" ) def is_base_parameter(name: str) -> bool: return name == "layer_logits" or name.startswith( ( "layer_fusion_net.", "local_projections.", "net.", "transition_head.", ) ) source = payload["head"] current = head.state_dict() required = [name for name in current if is_base_parameter(name)] missing = [name for name in required if name not in source] bad_shapes = [ name for name in required if name in source and source[name].shape != current[name].shape ] if missing or bad_shapes: raise ValueError( "Pooled checkpoint cannot initialize the exact base; " f"missing={missing[:8]}, shape_mismatch={bad_shapes[:8]}" ) for name in required: current[name] = source[name] head.load_state_dict(current) digest = hashlib.sha256() with path.open("rb") as stream: for chunk in iter(lambda: stream.read(1024 * 1024), b""): digest.update(chunk) args.initialize_pooled_checkpoint_sha256 = digest.hexdigest() args.initialize_pooled_checkpoint_epoch = payload.get("epoch") if not args.train_pooled_base: for name, parameter in head.named_parameters(): if is_base_parameter(name): parameter.requires_grad_(False) print( f"initialized pooled base: {path}; " f"frozen={not args.train_pooled_base}", flush=True, ) def main() -> None: args = parse_args() _validate_args(args) if args.layer_fusion is None: args.layer_fusion = ( "scalar-mixture" if args.learnable_layer_mixture else "single" ) args.learnable_layer_mixture = args.layer_fusion == "scalar-mixture" _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, transition_targets, hidden_size, num_input_layers, feature_metadata, local_features, person_features, ) = _load_examples( args.feature_root, args.target_root, train_sessions, expected_mode, args.dataset, sparse_anchor_indices, args.local_feature_root, args.local_modalities, compute_transition_targets=args.transition_objective != "none", local_fusion=args.local_fusion, person_track_max_gap_seconds=args.person_track_max_gap_seconds, ) local_contract = feature_metadata[0].get("local_contract") if not len(states): raise ValueError("No valid training examples") print( f"training examples: {len(states)}, feature_dim={hidden_size}, " f"input_layers={num_input_layers}" ) if args.transition_objective == "none": print("transition targets: disabled (not derived or loaded)") else: print( "transition targets: " f"{int(transition_targets.sum().item())}/" f"{transition_targets.numel()} positive " f"({float(transition_targets.mean()) * 100:.3f}%)" ) if num_input_layers > 1 and args.layer_fusion == "single": raise ValueError( f"Features contain {num_input_layers} stacked layers; pass " "--layer-fusion fixed-mean, scalar-mixture, or concat-mlp" ) if num_input_layers == 1 and args.layer_fusion != "single": raise ValueError( f"--layer-fusion {args.layer_fusion} requires stacked [T,L,D] " "features" ) dev_states: torch.Tensor | None = None dev_targets: torch.Tensor | None = None dev_local: torch.Tensor | None = None dev_people: FramePersonFeatures | None = None if dev_sessions: ( dev_states, dev_targets, _, dev_hidden_size, dev_num_input_layers, _, dev_local, dev_people, ) = _load_examples( args.feature_root, args.target_root, dev_sessions, expected_mode, args.dataset, None, args.local_feature_root, args.local_modalities, local_contract, compute_transition_targets=args.transition_objective != "none", local_fusion=args.local_fusion, person_track_max_gap_seconds=args.person_track_max_gap_seconds, ) if dev_hidden_size != hidden_size: raise ValueError( f"Dev feature dimension {dev_hidden_size} differs from " f"training dimension {hidden_size}" ) if dev_num_input_layers != num_input_layers: raise ValueError( f"Dev feature layer count {dev_num_input_layers} differs from " f"training count {num_input_layers}" ) print(f"development examples: {len(dev_states)}") dataset = TensorDataset( states, targets, transition_targets, local_features, torch.arange(len(states), dtype=torch.long), ) 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, ) num_output_classes = output_classes_for_loss(args.loss) from local_turn_features import local_head_dimensions head = TurnPredictionHead( hidden_size=hidden_size, num_horizons=NUM_HORIZONS, num_classes=num_output_classes, dropout=args.dropout, num_input_layers=num_input_layers, learnable_layer_mixture=args.learnable_layer_mixture, layer_fusion=args.layer_fusion, transition_auxiliary=args.transition_objective != "none", local_projection_dim=args.local_projection_dim, local_fusion=args.local_fusion, person_attention_dim=args.person_attention_dim, person_attention_heads=args.person_attention_heads, person_history_frames=args.person_history_frames, person_track_max_gap_seconds=args.person_track_max_gap_seconds, **local_head_dimensions(local_contract), ).to(device) if args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES: assert args.initialize_pooled_checkpoint is not None _initialize_person_residual_from_pooled( head, args.initialize_pooled_checkpoint, args=args, train_sessions=train_sessions, dev_sessions=dev_sessions, local_contract=local_contract, ) if local_contract is not None: print( f"local fusion: {args.local_fusion}/{args.local_modalities}; " f"raw V={head.local_visual_dim}, A={head.local_audio_dim}; " f"projection={args.local_projection_dim or 'identity'}" ) if person_features is not None: print( "frame sets: " f"{len(person_features.frame_offsets) - 1} source frames, " f"{len(person_features.features)} person observations, " f"max_people={person_features.max_people}, " f"attention={args.person_attention_dim}D/" f"{args.person_attention_heads} heads" ) trainable_parameters = sum( parameter.numel() for parameter in head.parameters() if parameter.requires_grad ) fusion_parameters = ( sum(parameter.numel() for parameter in head.layer_fusion_net.parameters()) if head.layer_fusion_net is not None else (head.layer_logits.numel() if head.layer_logits is not None else 0) ) print( f"loss={args.loss}, output_space={output_space_for_loss(args.loss)}, " f"transition_objective={args.transition_objective}, " f"activity_outputs={num_output_classes}, " "transition_outputs=" f"{NUM_HORIZONS if args.transition_objective != 'none' else 0}, " f"layer_fusion={args.layer_fusion}, " f"fusion_parameters={fusion_parameters:,}, " f"head_trainable_parameters={trainable_parameters:,}" ) print( "local_fusion_scope=" f"{LOCAL_FUSION_SCOPES[args.local_fusion]}" ) if head.layer_logits is None: optimizer_parameters: Any = head.parameters() else: optimizer_parameters = [ { "params": [ parameter for name, parameter in head.named_parameters() if name != "layer_logits" ], "weight_decay": args.weight_decay, }, { "params": [head.layer_logits], # Weight decay on softmax logits artificially pulls the layer # distribution back toward uniform. "weight_decay": 0.0, }, ] optimizer = torch.optim.Adam( optimizer_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 = { "lstr-mce": egospeak_multihorizon_loss, "miniroad-normalized": miniroad_multihorizon_loss, "independent-bce": independent_bce_multihorizon_loss, "four-state-ce": four_state_multihorizon_loss, }[args.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() if ( args.local_fusion in INITIALIZED_PERSON_RESIDUAL_MODES and not args.train_pooled_base ): # Keep dropout/norm behavior in the frozen pooled checkpoint fixed # while the new person branch learns a residual correction. This # applies to both track-history adapters and the high-rate person # set; otherwise head.train() would silently turn dropout back on # inside the supposedly frozen matched base. if head.layer_fusion_net is not None: head.layer_fusion_net.eval() head.local_projections.eval() head.net.eval() if head.transition_head is not None: head.transition_head.eval() loss_sum = 0.0 activity_loss_sum = 0.0 transition_loss_sum = 0.0 transition_reward_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, transition_target_batch, local_batch, example_index_batch, ) in progress: local_batch = local_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) 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, ) transition_target_batch = transition_target_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) person_batch = None person_mask = None if person_features is not None: if head.local_fusion in TRACK_PERSON_FUSION_MODES: person_batch, person_mask = person_features.padded_history( example_index_batch, max_observations=head.person_history_frames, ) else: person_batch, person_mask = person_features.padded( example_index_batch ) person_batch = person_batch.to( device=device, dtype=torch.float32, non_blocking=True, ) person_mask = person_mask.to(device=device, non_blocking=True) optimizer.zero_grad(set_to_none=True) if args.transition_objective == "none": logits = head( state_batch, local_batch, person_batch, person_mask, ) activity_loss = loss_function(logits, target_batch) loss = activity_loss transition_loss = activity_loss.detach().new_zeros(()) expected_reward = activity_loss.detach().new_zeros(()) else: logits, transition_logits = head.forward_with_transition( state_batch, local_batch, person_batch, person_mask, ) activity_loss = loss_function(logits, target_batch) components = transition_aware_objective( activity_loss, transition_logits, transition_target_batch, objective=args.transition_objective, transition_weight=args.transition_loss_weight, pos_weight=args.transition_pos_weight, reward_weight=args.transition_reward_weight, ) loss = components["total"] transition_loss = components["transition"] expected_reward = components["expected_reward"] loss.backward() optimizer.step() scheduler.step() batch_count = len(state_batch) loss_sum += float(loss.detach()) * batch_count activity_loss_sum += float(activity_loss.detach()) * batch_count transition_loss_sum += float(transition_loss.detach()) * batch_count transition_reward_sum += float(expected_reward.detach()) * batch_count example_count += batch_count postfix = { "loss": f"{float(loss.detach()):.5f}", "lr": f"{scheduler.get_last_lr()[0]:.3g}", } if args.transition_objective != "none": postfix["transition"] = f"{float(transition_loss.detach()):.5f}" postfix["reward"] = f"{float(expected_reward.detach()):.5f}" progress.set_postfix(postfix) train_loss = loss_sum / example_count record: dict[str, Any] = { "epoch": epoch, "train_loss": train_loss, "train_activity_loss": activity_loss_sum / example_count, "learning_rate": scheduler.get_last_lr()[0], } if args.transition_objective != "none": record["train_transition_loss"] = ( transition_loss_sum / example_count ) record["train_expected_transition_reward"] = ( transition_reward_sum / example_count ) layer_weights = head.layer_mixture_weights() if layer_weights is not None: record["layer_mixture_weights_oldest_to_newest"] = ( layer_weights.detach().float().cpu().tolist() ) interaction_gates = head.person_interaction_gates() if interaction_gates: record["person_interaction_gates"] = interaction_gates adapter_diagnostics = head.person_adapter_diagnostics() if adapter_diagnostics: record["person_adapter_diagnostics"] = adapter_diagnostics weight_text = ( "" if layer_weights is None else ", layer_weights=[" + ",".join( f"{value:.4f}" for value in record["layer_mixture_weights_oldest_to_newest"] ) + "]" ) transition_text = ( "" if args.transition_objective == "none" else ( f", activity_loss={record['train_activity_loss']:.6f}, " f"transition_loss={record['train_transition_loss']:.6f}, " "expected_transition_reward=" f"{record['train_expected_transition_reward']:.6f}" ) ) gate_text = ( "" if not interaction_gates else ", person_gates=" + json.dumps(interaction_gates, sort_keys=True) ) adapter_text = ( "" if not adapter_diagnostics else ", person_adapter=" + json.dumps(adapter_diagnostics, sort_keys=True) ) 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, args.loss, dev_local, dev_people, ) 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}" f"{transition_text}{weight_text}{gate_text}{adapter_text}" ) 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, num_input_layers=num_input_layers, 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}" f"{transition_text}{weight_text}{gate_text}{adapter_text}" ) 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, num_input_layers=num_input_layers, 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 payload["head_config"].get("layer_fusion") == "scalar-mixture": saved_logits = payload["head"]["layer_logits"].float() saved_weights = saved_logits.softmax(dim=0).tolist() print( "saved layer mixture weights (oldest->newest): " + ", ".join(f"{value:.6f}" for value in saved_weights) ) if __name__ == "__main__": main()