1272 lines
48 KiB
Python
1272 lines
48 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,
|
|
LOCAL_FUSION_MODES,
|
|
LOCAL_FUSION_SCOPES,
|
|
LOSS_NAMES,
|
|
NUM_HORIZONS,
|
|
OFFICIAL_TRAIN_SESSIONS,
|
|
STEP_SECONDS,
|
|
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
|
|
|
|
|
|
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 adds a zero-gated causal track correction; "
|
|
"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 track-gru-residual: copy the exact pooled G+A+V base "
|
|
"before training the zero-gated track 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 == "track-gru-residual":
|
|
if args.initialize_pooled_checkpoint is None:
|
|
raise ValueError(
|
|
"track-gru-residual 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 "
|
|
"--local-fusion track-gru-residual"
|
|
)
|
|
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 == "track-gru-residual":
|
|
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
|
|
),
|
|
"train_pooled_base": 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_track_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 inside a track 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(
|
|
"Track 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 track experiment")
|
|
if payload.get("dataset", args.dataset) != args.dataset:
|
|
raise ValueError("Pooled checkpoint dataset differs from the track 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 payload.get("local_contract") != local_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": head.local_visual_dim,
|
|
"local_audio_dim": head.local_audio_dim,
|
|
"local_projection_dim": head.local_projection_dim,
|
|
}
|
|
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 track base: {mismatches}"
|
|
)
|
|
|
|
def is_base_parameter(name: str) -> bool:
|
|
return name == "layer_logits" or name.startswith(
|
|
("layer_fusion_net.", "local_projections.", "net.")
|
|
)
|
|
|
|
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)
|
|
|
|
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 == "track-gru-residual":
|
|
assert args.initialize_pooled_checkpoint is not None
|
|
_initialize_track_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 == "track-gru-residual" and not args.train_pooled_base:
|
|
# Keep dropout/norm behavior in the frozen pooled checkpoint fixed
|
|
# while the new track branch learns a residual correction.
|
|
if head.layer_fusion_net is not None:
|
|
head.layer_fusion_net.eval()
|
|
head.local_projections.eval()
|
|
head.net.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 == "track-gru-residual":
|
|
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
|
|
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)
|
|
)
|
|
|
|
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}"
|
|
)
|
|
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}"
|
|
)
|
|
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()
|