954 lines
35 KiB
Python
954 lines
35 KiB
Python
"""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,
|
|
LAYER_FUSION_MODES,
|
|
LOSS_NAMES,
|
|
NUM_HORIZONS,
|
|
OFFICIAL_TRAIN_SESSIONS,
|
|
STEP_SECONDS,
|
|
TRANSITION_OBJECTIVE_NAMES,
|
|
TurnPredictionHead,
|
|
average_precision_metrics,
|
|
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,
|
|
)
|
|
|
|
|
|
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-projection-dim", type=int, default=256,
|
|
help="Trainable projection per local modality; 0 retains encoder dimensions.",
|
|
)
|
|
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.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,
|
|
) -> tuple[
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
torch.Tensor,
|
|
int,
|
|
int,
|
|
list[dict[str, Any]],
|
|
torch.Tensor,
|
|
]:
|
|
all_states: list[torch.Tensor] = []
|
|
all_targets: list[torch.Tensor] = []
|
|
all_transition_targets: list[torch.Tensor] = []
|
|
all_local: list[torch.Tensor] = []
|
|
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_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)
|
|
)
|
|
activity_sequence = torch.cat(
|
|
(data.targets[current_indices, None, :], targets),
|
|
dim=1,
|
|
)
|
|
transition_targets = transition_targets_from_activity(activity_sequence)
|
|
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 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)
|
|
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")
|
|
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),
|
|
)
|
|
|
|
|
|
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,
|
|
) -> 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
|
|
)
|
|
chunks.append(marginal_probabilities_from_logits(head(x, side), 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": 4,
|
|
"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_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_projection_dim": args.local_projection_dim,
|
|
"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 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,
|
|
) = _load_examples(
|
|
args.feature_root,
|
|
args.target_root,
|
|
train_sessions,
|
|
expected_mode,
|
|
args.dataset,
|
|
sparse_anchor_indices,
|
|
args.local_feature_root,
|
|
args.local_modalities,
|
|
)
|
|
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}"
|
|
)
|
|
print(
|
|
"transition targets: "
|
|
f"{int(transition_targets.sum().item())}/{transition_targets.numel()} "
|
|
f"positive ({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
|
|
if dev_sessions:
|
|
(
|
|
dev_states,
|
|
dev_targets,
|
|
_,
|
|
dev_hidden_size,
|
|
dev_num_input_layers,
|
|
_,
|
|
dev_local,
|
|
) = _load_examples(
|
|
args.feature_root,
|
|
args.target_root,
|
|
dev_sessions,
|
|
expected_mode,
|
|
args.dataset,
|
|
None,
|
|
args.local_feature_root,
|
|
args.local_modalities,
|
|
local_contract,
|
|
)
|
|
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)
|
|
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_head_dimensions(local_contract),
|
|
).to(device)
|
|
if local_contract is not None:
|
|
print(
|
|
f"local fusion: {args.local_modalities}; raw V={head.local_visual_dim}, "
|
|
f"A={head.local_audio_dim}; projection={args.local_projection_dim or 'identity'}"
|
|
)
|
|
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:,}"
|
|
)
|
|
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()
|
|
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 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,
|
|
)
|
|
optimizer.zero_grad(set_to_none=True)
|
|
if args.transition_objective == "none":
|
|
logits = head(state_batch, local_batch)
|
|
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)
|
|
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()
|
|
)
|
|
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}"
|
|
)
|
|
)
|
|
|
|
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_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}"
|
|
)
|
|
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}"
|
|
)
|
|
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()
|