632 lines
22 KiB
Python
632 lines
22 KiB
Python
|
#!/usr/bin/env python3
|
||
|
|
"""Train stage-one frozen-EgoSpeak plus frozen-Omni residual fusion.
|
||
|
|
|
||
|
|
The released MiniROADA A+V checkpoint is loaded strictly and remains in eval
|
||
|
|
mode with gradients disabled. Existing causal 5 Hz Omni states are fused with
|
||
|
|
the MiniROADA GRU state, and only a zero-initialized residual anticipation head
|
||
|
|
is optimized. Evaluation happens once after the requested fixed epoch budget;
|
||
|
|
the script saves both fused and exactly matched frozen-base predictions.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import json
|
||
|
|
import math
|
||
|
|
import random
|
||
|
|
from collections.abc import Sequence
|
||
|
|
from contextlib import nullcontext
|
||
|
|
from pathlib import Path
|
||
|
|
from typing import Any
|
||
|
|
|
||
|
|
import numpy as np
|
||
|
|
import torch
|
||
|
|
from torch.utils.data import DataLoader
|
||
|
|
from tqdm.auto import tqdm
|
||
|
|
|
||
|
|
from easycom_head import (
|
||
|
|
CLASS_NAMES,
|
||
|
|
NUM_HORIZONS,
|
||
|
|
OFFICIAL_TEST_SESSIONS,
|
||
|
|
OFFICIAL_TRAIN_SESSIONS,
|
||
|
|
STEP_SECONDS,
|
||
|
|
average_precision_metrics,
|
||
|
|
)
|
||
|
|
from egospeak_omni_residual import (
|
||
|
|
EasyComSessionArrays,
|
||
|
|
EgoSpeakOmniResidual,
|
||
|
|
EgoSpeakOmniWindowDataset,
|
||
|
|
align_source_predictions,
|
||
|
|
load_easycom_session_arrays,
|
||
|
|
load_frozen_miniroada,
|
||
|
|
miniroad_normalized_loss,
|
||
|
|
normalize_session_name,
|
||
|
|
sha256_file,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args() -> argparse.Namespace:
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument(
|
||
|
|
"--egospeak-feature-root",
|
||
|
|
type=Path,
|
||
|
|
required=True,
|
||
|
|
help=(
|
||
|
|
"EasyCom root containing rgb_kinetics_resnet50/, "
|
||
|
|
"audio_wav2vec2/, and target_perframe/."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--omni-feature-root",
|
||
|
|
type=Path,
|
||
|
|
required=True,
|
||
|
|
help="Existing directory containing causal Session_X.pt Omni states.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--base-checkpoint",
|
||
|
|
type=Path,
|
||
|
|
required=True,
|
||
|
|
help="Released EgoSpeak miniroad_easycom_AV.pth state dict.",
|
||
|
|
)
|
||
|
|
parser.add_argument("--output-dir", type=Path, required=True)
|
||
|
|
parser.add_argument("--train-sessions", nargs="+", default=None)
|
||
|
|
parser.add_argument("--eval-sessions", nargs="+", default=None)
|
||
|
|
parser.add_argument(
|
||
|
|
"--expected-omni-mode",
|
||
|
|
default="open-stream-replay",
|
||
|
|
help="Expected saved extraction mode; pass 'any' only for an audit run.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--omni-layer-fusion",
|
||
|
|
choices=("single", "last", "mean"),
|
||
|
|
default="mean",
|
||
|
|
help="How to collapse saved [T,L,D] Thinker states.",
|
||
|
|
)
|
||
|
|
parser.add_argument("--epochs", type=int, default=4)
|
||
|
|
parser.add_argument("--batch-size", type=int, default=64)
|
||
|
|
parser.add_argument(
|
||
|
|
"--eval-chunk-size",
|
||
|
|
type=int,
|
||
|
|
default=0,
|
||
|
|
help=(
|
||
|
|
"GRU inference chunk size; 0 processes each session in one pass "
|
||
|
|
"to match released MiniROADA evaluation."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
parser.add_argument("--num-workers", type=int, default=4)
|
||
|
|
parser.add_argument("--window-size", type=int, default=128)
|
||
|
|
parser.add_argument("--stride", type=int, default=4)
|
||
|
|
parser.add_argument("--fusion-dim", type=int, default=256)
|
||
|
|
parser.add_argument("--tail-embedding-dim", type=int, default=32)
|
||
|
|
parser.add_argument("--dropout", type=float, default=0.2)
|
||
|
|
parser.add_argument("--lr", type=float, default=1e-4)
|
||
|
|
parser.add_argument("--weight-decay", type=float, default=0.05)
|
||
|
|
parser.add_argument("--seed", type=int, default=0)
|
||
|
|
parser.add_argument("--device", default="cuda")
|
||
|
|
parser.add_argument("--amp", action="store_true")
|
||
|
|
parser.add_argument(
|
||
|
|
"--max-train-examples",
|
||
|
|
type=int,
|
||
|
|
help="Smoke-test cap; omit for a complete run.",
|
||
|
|
)
|
||
|
|
parser.add_argument(
|
||
|
|
"--reference-prediction-root",
|
||
|
|
type=Path,
|
||
|
|
help=(
|
||
|
|
"Optional released MiniROADA raw eval directory. The frozen base "
|
||
|
|
"is compared against Session_X.npy files before results are accepted."
|
||
|
|
),
|
||
|
|
)
|
||
|
|
parser.add_argument("--reference-tolerance", type=float, default=1e-5)
|
||
|
|
parser.add_argument("--no-progress", action="store_true")
|
||
|
|
return parser.parse_args()
|
||
|
|
|
||
|
|
|
||
|
|
def set_seed(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)
|
||
|
|
torch.backends.cudnn.benchmark = False
|
||
|
|
torch.backends.cudnn.deterministic = True
|
||
|
|
|
||
|
|
|
||
|
|
def resolve_device(value: str) -> torch.device:
|
||
|
|
device = torch.device(value)
|
||
|
|
if device.type == "cuda" and not torch.cuda.is_available():
|
||
|
|
raise RuntimeError("CUDA was requested but is unavailable")
|
||
|
|
return device
|
||
|
|
|
||
|
|
|
||
|
|
def normalize_sessions(
|
||
|
|
values: Sequence[str] | None,
|
||
|
|
default: Sequence[str],
|
||
|
|
) -> list[str]:
|
||
|
|
sessions = [normalize_session_name(value) for value in (values or default)]
|
||
|
|
if len(sessions) != len(set(sessions)):
|
||
|
|
raise ValueError("Duplicate sessions were supplied")
|
||
|
|
return sessions
|
||
|
|
|
||
|
|
|
||
|
|
def load_sessions(
|
||
|
|
args: argparse.Namespace,
|
||
|
|
sessions: Sequence[str],
|
||
|
|
) -> list[EasyComSessionArrays]:
|
||
|
|
expected_mode = (
|
||
|
|
None if args.expected_omni_mode == "any" else args.expected_omni_mode
|
||
|
|
)
|
||
|
|
loaded = []
|
||
|
|
for session in sessions:
|
||
|
|
item = load_easycom_session_arrays(
|
||
|
|
args.egospeak_feature_root,
|
||
|
|
args.omni_feature_root,
|
||
|
|
session,
|
||
|
|
expected_omni_mode=expected_mode,
|
||
|
|
omni_layer_fusion=args.omni_layer_fusion,
|
||
|
|
)
|
||
|
|
loaded.append(item)
|
||
|
|
print(
|
||
|
|
f"{item.name}: steps={item.num_steps}, "
|
||
|
|
f"omni={tuple(item.omni.shape)}, "
|
||
|
|
f"mode={item.omni_metadata['mode']}"
|
||
|
|
)
|
||
|
|
dims = {int(item.omni.shape[-1]) for item in loaded}
|
||
|
|
if len(dims) != 1:
|
||
|
|
raise ValueError(f"Omni dimensions differ across sessions: {sorted(dims)}")
|
||
|
|
signatures = {
|
||
|
|
(
|
||
|
|
item.omni_metadata["model_id"],
|
||
|
|
item.omni_metadata["mode"],
|
||
|
|
item.omni_metadata["config"].get("memory_seconds"),
|
||
|
|
item.omni_metadata["config"].get("video_fps"),
|
||
|
|
item.omni_metadata["config"].get("readout_last_n_layers"),
|
||
|
|
item.omni_metadata["config"].get("readout_layer_fusion"),
|
||
|
|
tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]),
|
||
|
|
)
|
||
|
|
for item in loaded
|
||
|
|
}
|
||
|
|
if len(signatures) != 1:
|
||
|
|
raise ValueError(
|
||
|
|
"Selected sessions mix incompatible Omni extraction contracts: "
|
||
|
|
f"{sorted(map(str, signatures))}"
|
||
|
|
)
|
||
|
|
return loaded
|
||
|
|
|
||
|
|
|
||
|
|
def omni_contract_signature(item: EasyComSessionArrays) -> tuple[Any, ...]:
|
||
|
|
config = item.omni_metadata["config"]
|
||
|
|
return (
|
||
|
|
item.omni_metadata["model_id"],
|
||
|
|
item.omni_metadata["mode"],
|
||
|
|
config.get("memory_seconds"),
|
||
|
|
config.get("video_fps"),
|
||
|
|
config.get("readout_last_n_layers"),
|
||
|
|
config.get("readout_layer_fusion"),
|
||
|
|
tuple(item.omni_metadata["state_shape_before_layer_fusion"][1:]),
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
def autocast_context(device: torch.device, enabled: bool):
|
||
|
|
if not enabled:
|
||
|
|
return nullcontext()
|
||
|
|
if device.type != "cuda":
|
||
|
|
raise ValueError("--amp currently requires CUDA")
|
||
|
|
return torch.autocast(device_type="cuda", dtype=torch.float16)
|
||
|
|
|
||
|
|
|
||
|
|
def make_loader(
|
||
|
|
dataset: EgoSpeakOmniWindowDataset,
|
||
|
|
args: argparse.Namespace,
|
||
|
|
epoch: int,
|
||
|
|
) -> DataLoader:
|
||
|
|
generator = torch.Generator()
|
||
|
|
generator.manual_seed(args.seed + epoch)
|
||
|
|
return DataLoader(
|
||
|
|
dataset,
|
||
|
|
batch_size=args.batch_size,
|
||
|
|
shuffle=True,
|
||
|
|
num_workers=args.num_workers,
|
||
|
|
pin_memory=torch.cuda.is_available(),
|
||
|
|
persistent_workers=False,
|
||
|
|
generator=generator,
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
@torch.inference_mode()
|
||
|
|
def verify_zero_initialization(
|
||
|
|
model: EgoSpeakOmniResidual,
|
||
|
|
dataset: EgoSpeakOmniWindowDataset,
|
||
|
|
device: torch.device,
|
||
|
|
) -> float:
|
||
|
|
model.eval()
|
||
|
|
rgb, audio, omni, tail_bin, _, _, _ = dataset[0]
|
||
|
|
fused, base = model.forward_last(
|
||
|
|
rgb[None].to(device),
|
||
|
|
audio[None].to(device),
|
||
|
|
omni[None].to(device),
|
||
|
|
tail_bin[None].to(device),
|
||
|
|
)
|
||
|
|
maximum = float((fused - base).abs().max().cpu())
|
||
|
|
if maximum != 0.0:
|
||
|
|
raise RuntimeError(
|
||
|
|
"Zero-initialized residual does not exactly recover MiniROADA: "
|
||
|
|
f"max_abs_difference={maximum}"
|
||
|
|
)
|
||
|
|
print("Initialization audit: fused logits == frozen EgoSpeak logits exactly")
|
||
|
|
return maximum
|
||
|
|
|
||
|
|
|
||
|
|
def train_one_epoch(
|
||
|
|
model: EgoSpeakOmniResidual,
|
||
|
|
loader: DataLoader,
|
||
|
|
optimizer: torch.optim.Optimizer,
|
||
|
|
scaler: torch.cuda.amp.GradScaler,
|
||
|
|
device: torch.device,
|
||
|
|
*,
|
||
|
|
amp: bool,
|
||
|
|
show_progress: bool,
|
||
|
|
) -> dict[str, float]:
|
||
|
|
model.train()
|
||
|
|
total_loss = 0.0
|
||
|
|
total_examples = 0
|
||
|
|
iterator = tqdm(loader, desc="train", disable=not show_progress, dynamic_ncols=True)
|
||
|
|
for rgb, audio, omni, tail_bin, targets, _, _ in iterator:
|
||
|
|
rgb = rgb.to(device=device, dtype=torch.float32, non_blocking=True)
|
||
|
|
audio = audio.to(device=device, dtype=torch.float32, non_blocking=True)
|
||
|
|
omni = omni.to(device=device, dtype=torch.float32, non_blocking=True)
|
||
|
|
tail_bin = tail_bin.to(device=device, dtype=torch.long, non_blocking=True)
|
||
|
|
targets = targets.to(device=device, dtype=torch.float32, non_blocking=True)
|
||
|
|
|
||
|
|
optimizer.zero_grad(set_to_none=True)
|
||
|
|
with autocast_context(device, amp):
|
||
|
|
logits, _ = model.forward_last(rgb, audio, omni, tail_bin)
|
||
|
|
loss = miniroad_normalized_loss(logits.float(), targets.float())
|
||
|
|
scaler.scale(loss).backward()
|
||
|
|
scaler.step(optimizer)
|
||
|
|
scaler.update()
|
||
|
|
|
||
|
|
batch_size = len(rgb)
|
||
|
|
total_examples += batch_size
|
||
|
|
total_loss += float(loss.detach().cpu()) * batch_size
|
||
|
|
iterator.set_postfix_str(f"loss={total_loss / total_examples:.5f}")
|
||
|
|
return {
|
||
|
|
"loss": total_loss / total_examples,
|
||
|
|
"num_examples": total_examples,
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
@torch.inference_mode()
|
||
|
|
def predict_session(
|
||
|
|
model: EgoSpeakOmniResidual,
|
||
|
|
session: EasyComSessionArrays,
|
||
|
|
device: torch.device,
|
||
|
|
*,
|
||
|
|
chunk_size: int,
|
||
|
|
amp: bool,
|
||
|
|
show_progress: bool,
|
||
|
|
) -> tuple[np.ndarray, np.ndarray]:
|
||
|
|
model.eval()
|
||
|
|
num_sources = session.num_steps - NUM_HORIZONS
|
||
|
|
if num_sources <= 0:
|
||
|
|
raise ValueError(f"{session.name}: no valid ten-horizon source rows")
|
||
|
|
if chunk_size == 0:
|
||
|
|
chunk_size = num_sources
|
||
|
|
fused_chunks: list[torch.Tensor] = []
|
||
|
|
base_chunks: list[torch.Tensor] = []
|
||
|
|
recurrent_state = None
|
||
|
|
starts = range(0, num_sources, chunk_size)
|
||
|
|
for start in tqdm(
|
||
|
|
starts,
|
||
|
|
total=math.ceil(num_sources / chunk_size),
|
||
|
|
desc=f"infer {session.name}",
|
||
|
|
disable=not show_progress,
|
||
|
|
dynamic_ncols=True,
|
||
|
|
):
|
||
|
|
end = min(start + chunk_size, num_sources)
|
||
|
|
rgb = torch.from_numpy(
|
||
|
|
np.asarray(session.rgb[start:end], dtype=np.float32).copy()
|
||
|
|
)[None].to(device)
|
||
|
|
audio = torch.from_numpy(
|
||
|
|
np.asarray(session.audio[start:end], dtype=np.float32).copy()
|
||
|
|
)[None].to(device)
|
||
|
|
omni = session.omni[start:end][None].to(device=device, dtype=torch.float32)
|
||
|
|
tail_bin = session.tail_bins[start:end][None].to(device)
|
||
|
|
with autocast_context(device, amp):
|
||
|
|
fused, base, recurrent_state = model.forward_sequence_chunk(
|
||
|
|
rgb,
|
||
|
|
audio,
|
||
|
|
omni,
|
||
|
|
tail_bin,
|
||
|
|
recurrent_state,
|
||
|
|
)
|
||
|
|
recurrent_state = recurrent_state.detach()
|
||
|
|
fused_chunks.append(torch.softmax(fused.float(), dim=-1)[0].cpu())
|
||
|
|
base_chunks.append(torch.softmax(base.float(), dim=-1)[0].cpu())
|
||
|
|
return torch.cat(fused_chunks).numpy(), torch.cat(base_chunks).numpy()
|
||
|
|
|
||
|
|
|
||
|
|
def evaluate_predictions(
|
||
|
|
predictions: dict[str, np.ndarray],
|
||
|
|
valid: dict[str, np.ndarray],
|
||
|
|
targets: dict[str, np.ndarray],
|
||
|
|
) -> dict[str, Any]:
|
||
|
|
full_results = []
|
||
|
|
causal_results = []
|
||
|
|
for horizon in range(NUM_HORIZONS):
|
||
|
|
full_scores = np.concatenate(
|
||
|
|
[item[:, :, horizon] for item in predictions.values()], axis=0
|
||
|
|
)
|
||
|
|
full_targets = np.concatenate(list(targets.values()), axis=0)
|
||
|
|
full = average_precision_metrics(full_scores, full_targets)
|
||
|
|
full.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS)
|
||
|
|
full_results.append(full)
|
||
|
|
|
||
|
|
causal_scores = np.concatenate(
|
||
|
|
[
|
||
|
|
predictions[name][valid[name][:, horizon], :, horizon]
|
||
|
|
for name in predictions
|
||
|
|
],
|
||
|
|
axis=0,
|
||
|
|
)
|
||
|
|
causal_targets = np.concatenate(
|
||
|
|
[targets[name][valid[name][:, horizon]] for name in predictions],
|
||
|
|
axis=0,
|
||
|
|
)
|
||
|
|
causal = average_precision_metrics(causal_scores, causal_targets)
|
||
|
|
causal.update(horizon=horizon + 1, seconds=(horizon + 1) * STEP_SECONDS)
|
||
|
|
causal_results.append(causal)
|
||
|
|
return {
|
||
|
|
"egospeak_full_timeline": {
|
||
|
|
"horizons": full_results,
|
||
|
|
"mean_map": float(np.mean([item["map"] for item in full_results])),
|
||
|
|
},
|
||
|
|
"causal_valid_only": {
|
||
|
|
"horizons": causal_results,
|
||
|
|
"mean_map": float(np.mean([item["map"] for item in causal_results])),
|
||
|
|
},
|
||
|
|
}
|
||
|
|
|
||
|
|
|
||
|
|
def print_metrics(label: str, metrics: dict[str, Any]) -> None:
|
||
|
|
print(f"\n=== {label} ===")
|
||
|
|
for protocol_name in ("egospeak_full_timeline", "causal_valid_only"):
|
||
|
|
protocol = metrics[protocol_name]
|
||
|
|
print(f"{protocol_name}:")
|
||
|
|
for item in protocol["horizons"]:
|
||
|
|
class_text = ", ".join(
|
||
|
|
f"{name}={value * 100:.3f}" for name, value in item["class_ap"].items()
|
||
|
|
)
|
||
|
|
print(f" {item['seconds']:.1f}s mAP={item['map'] * 100:.3f}; {class_text}")
|
||
|
|
print(f" mean mAP={protocol['mean_map'] * 100:.3f}")
|
||
|
|
|
||
|
|
|
||
|
|
def prepare_output(path: Path) -> Path:
|
||
|
|
path = path.expanduser().resolve()
|
||
|
|
if path.exists() and any(path.iterdir()):
|
||
|
|
raise FileExistsError(f"Refusing to overwrite nonempty output: {path}")
|
||
|
|
path.mkdir(parents=True, exist_ok=True)
|
||
|
|
for name in (
|
||
|
|
"raw_eval",
|
||
|
|
"pred",
|
||
|
|
"valid",
|
||
|
|
"frozen_egospeak_raw_eval",
|
||
|
|
"frozen_egospeak_pred",
|
||
|
|
):
|
||
|
|
(path / name).mkdir()
|
||
|
|
return path
|
||
|
|
|
||
|
|
|
||
|
|
def main() -> None:
|
||
|
|
args = parse_args()
|
||
|
|
if args.epochs <= 0 or args.batch_size <= 0 or args.eval_chunk_size < 0:
|
||
|
|
raise ValueError(
|
||
|
|
"Epoch and batch sizes must be positive; evaluation chunk size "
|
||
|
|
"must be nonnegative"
|
||
|
|
)
|
||
|
|
if args.reference_tolerance < 0:
|
||
|
|
raise ValueError("--reference-tolerance must be nonnegative")
|
||
|
|
set_seed(args.seed)
|
||
|
|
device = resolve_device(args.device)
|
||
|
|
|
||
|
|
train_names = normalize_sessions(args.train_sessions, OFFICIAL_TRAIN_SESSIONS)
|
||
|
|
eval_names = normalize_sessions(args.eval_sessions, OFFICIAL_TEST_SESSIONS)
|
||
|
|
train_sessions = load_sessions(args, train_names)
|
||
|
|
eval_sessions = load_sessions(args, eval_names)
|
||
|
|
omni_dims = {int(item.omni.shape[-1]) for item in (*train_sessions, *eval_sessions)}
|
||
|
|
if len(omni_dims) != 1:
|
||
|
|
raise ValueError(f"Omni dimensions differ: {sorted(omni_dims)}")
|
||
|
|
contracts = {
|
||
|
|
omni_contract_signature(item) for item in (*train_sessions, *eval_sessions)
|
||
|
|
}
|
||
|
|
if len(contracts) != 1:
|
||
|
|
raise ValueError(
|
||
|
|
"Train/evaluation sessions use incompatible Omni extraction "
|
||
|
|
f"contracts: {sorted(map(str, contracts))}"
|
||
|
|
)
|
||
|
|
omni_dim = omni_dims.pop()
|
||
|
|
|
||
|
|
base = load_frozen_miniroada(args.base_checkpoint)
|
||
|
|
model = EgoSpeakOmniResidual(
|
||
|
|
base,
|
||
|
|
omni_dim=omni_dim,
|
||
|
|
fusion_dim=args.fusion_dim,
|
||
|
|
tail_embedding_dim=args.tail_embedding_dim,
|
||
|
|
dropout=args.dropout,
|
||
|
|
).to(device)
|
||
|
|
trainable = model.trainable_parameters()
|
||
|
|
print(f"device: {device}")
|
||
|
|
print(f"train sessions: {', '.join(train_names)}")
|
||
|
|
print(f"evaluation sessions: {', '.join(eval_names)}")
|
||
|
|
print(
|
||
|
|
f"parameters: total={sum(p.numel() for p in model.parameters()):,}, "
|
||
|
|
f"trainable={sum(p.numel() for p in trainable):,}, "
|
||
|
|
"MiniROADA frozen=True"
|
||
|
|
)
|
||
|
|
|
||
|
|
dataset = EgoSpeakOmniWindowDataset(
|
||
|
|
train_sessions,
|
||
|
|
window_size=args.window_size,
|
||
|
|
stride=args.stride,
|
||
|
|
seed=args.seed,
|
||
|
|
max_examples=args.max_train_examples,
|
||
|
|
)
|
||
|
|
initialization_difference = verify_zero_initialization(model, dataset, device)
|
||
|
|
output_dir = prepare_output(args.output_dir)
|
||
|
|
optimizer = torch.optim.AdamW(
|
||
|
|
trainable,
|
||
|
|
lr=args.lr,
|
||
|
|
weight_decay=args.weight_decay,
|
||
|
|
)
|
||
|
|
scaler = torch.cuda.amp.GradScaler(enabled=args.amp)
|
||
|
|
history = []
|
||
|
|
for epoch in range(1, args.epochs + 1):
|
||
|
|
dataset.resample(epoch - 1)
|
||
|
|
loader = make_loader(dataset, args, epoch)
|
||
|
|
result = train_one_epoch(
|
||
|
|
model,
|
||
|
|
loader,
|
||
|
|
optimizer,
|
||
|
|
scaler,
|
||
|
|
device,
|
||
|
|
amp=args.amp,
|
||
|
|
show_progress=not args.no_progress,
|
||
|
|
)
|
||
|
|
result["epoch"] = epoch
|
||
|
|
history.append(result)
|
||
|
|
print(
|
||
|
|
f"epoch {epoch}/{args.epochs}: loss={result['loss']:.6f}, "
|
||
|
|
f"examples={result['num_examples']}"
|
||
|
|
)
|
||
|
|
|
||
|
|
checkpoint_path = output_dir / "checkpoint_last.pt"
|
||
|
|
checkpoint_payload = {
|
||
|
|
"format": "egospeak-omni-zero-residual-v1",
|
||
|
|
"model_state_dict": model.state_dict(),
|
||
|
|
"model_config": {
|
||
|
|
"omni_dim": omni_dim,
|
||
|
|
"fusion_dim": args.fusion_dim,
|
||
|
|
"tail_embedding_dim": args.tail_embedding_dim,
|
||
|
|
"num_tail_bins": NUM_HORIZONS,
|
||
|
|
"dropout": args.dropout,
|
||
|
|
},
|
||
|
|
"base_checkpoint": str(args.base_checkpoint.expanduser().resolve()),
|
||
|
|
"base_checkpoint_sha256": sha256_file(
|
||
|
|
args.base_checkpoint.expanduser().resolve()
|
||
|
|
),
|
||
|
|
"train_sessions": train_names,
|
||
|
|
"eval_sessions": eval_names,
|
||
|
|
"training_config": {
|
||
|
|
"epochs": args.epochs,
|
||
|
|
"batch_size": args.batch_size,
|
||
|
|
"window_size": args.window_size,
|
||
|
|
"stride": args.stride,
|
||
|
|
"lr": args.lr,
|
||
|
|
"weight_decay": args.weight_decay,
|
||
|
|
"seed": args.seed,
|
||
|
|
"amp": args.amp,
|
||
|
|
"omni_layer_fusion": args.omni_layer_fusion,
|
||
|
|
"expected_omni_mode": args.expected_omni_mode,
|
||
|
|
"max_train_examples": args.max_train_examples,
|
||
|
|
},
|
||
|
|
"initialization_max_abs_logit_difference": initialization_difference,
|
||
|
|
"history": history,
|
||
|
|
}
|
||
|
|
torch.save(checkpoint_payload, checkpoint_path)
|
||
|
|
print(f"Saved checkpoint: {checkpoint_path}")
|
||
|
|
|
||
|
|
fused_predictions: dict[str, np.ndarray] = {}
|
||
|
|
base_predictions: dict[str, np.ndarray] = {}
|
||
|
|
valid: dict[str, np.ndarray] = {}
|
||
|
|
targets: dict[str, np.ndarray] = {}
|
||
|
|
reference_audit: dict[str, Any] = {}
|
||
|
|
for session in eval_sessions:
|
||
|
|
fused_raw, base_raw = predict_session(
|
||
|
|
model,
|
||
|
|
session,
|
||
|
|
device,
|
||
|
|
chunk_size=args.eval_chunk_size,
|
||
|
|
amp=args.amp,
|
||
|
|
show_progress=not args.no_progress,
|
||
|
|
)
|
||
|
|
fused_aligned, session_valid = align_source_predictions(
|
||
|
|
fused_raw, session.num_steps
|
||
|
|
)
|
||
|
|
base_aligned, base_valid = align_source_predictions(base_raw, session.num_steps)
|
||
|
|
if not np.array_equal(session_valid, base_valid):
|
||
|
|
raise RuntimeError("Fused/base causal-valid masks unexpectedly differ")
|
||
|
|
np.save(output_dir / "raw_eval" / f"{session.name}.npy", fused_raw)
|
||
|
|
np.save(output_dir / "pred" / f"{session.name}.npy", fused_aligned)
|
||
|
|
np.save(output_dir / "valid" / f"{session.name}.npy", session_valid)
|
||
|
|
np.save(
|
||
|
|
output_dir / "frozen_egospeak_raw_eval" / f"{session.name}.npy",
|
||
|
|
base_raw,
|
||
|
|
)
|
||
|
|
np.save(
|
||
|
|
output_dir / "frozen_egospeak_pred" / f"{session.name}.npy",
|
||
|
|
base_aligned,
|
||
|
|
)
|
||
|
|
fused_predictions[session.name] = fused_aligned
|
||
|
|
base_predictions[session.name] = base_aligned
|
||
|
|
valid[session.name] = session_valid
|
||
|
|
targets[session.name] = np.asarray(session.targets, dtype=np.float32)
|
||
|
|
|
||
|
|
if args.reference_prediction_root is not None:
|
||
|
|
reference_path = (
|
||
|
|
args.reference_prediction_root.expanduser().resolve()
|
||
|
|
/ f"{session.name}.npy"
|
||
|
|
)
|
||
|
|
if not reference_path.is_file():
|
||
|
|
raise FileNotFoundError(reference_path)
|
||
|
|
reference = np.load(reference_path, allow_pickle=False)
|
||
|
|
if reference.shape != base_raw.shape:
|
||
|
|
raise ValueError(
|
||
|
|
f"{session.name}: reference shape {reference.shape} != "
|
||
|
|
f"{base_raw.shape}"
|
||
|
|
)
|
||
|
|
difference = np.abs(reference.astype(np.float32) - base_raw)
|
||
|
|
maximum = float(difference.max())
|
||
|
|
reference_audit[session.name] = {
|
||
|
|
"path": str(reference_path),
|
||
|
|
"max_abs_difference": maximum,
|
||
|
|
"mean_abs_difference": float(difference.mean()),
|
||
|
|
}
|
||
|
|
print(
|
||
|
|
f"{session.name}: frozen-base/reference max abs "
|
||
|
|
f"difference={maximum:.8g}"
|
||
|
|
)
|
||
|
|
if maximum > args.reference_tolerance:
|
||
|
|
raise RuntimeError(
|
||
|
|
f"{session.name}: frozen MiniROADA did not reproduce the "
|
||
|
|
f"reference within {args.reference_tolerance}: {maximum}"
|
||
|
|
)
|
||
|
|
|
||
|
|
fused_metrics = evaluate_predictions(fused_predictions, valid, targets)
|
||
|
|
base_metrics = evaluate_predictions(base_predictions, valid, targets)
|
||
|
|
metrics = {
|
||
|
|
**fused_metrics,
|
||
|
|
"format": "egospeak-omni-zero-residual-eval-v1",
|
||
|
|
"classes": list(CLASS_NAMES),
|
||
|
|
"sessions": eval_names,
|
||
|
|
"prediction_layout": "[target_time, class, future_horizon]",
|
||
|
|
"raw_prediction_layout": "[source_time, future_horizon, class]",
|
||
|
|
"checkpoint": str(checkpoint_path),
|
||
|
|
"train_sessions": train_names,
|
||
|
|
"frozen_egospeak_baseline": base_metrics,
|
||
|
|
"reference_prediction_audit": reference_audit,
|
||
|
|
"initialization_max_abs_logit_difference": initialization_difference,
|
||
|
|
"omni_sources": {
|
||
|
|
item.name: item.omni_metadata for item in (*train_sessions, *eval_sessions)
|
||
|
|
},
|
||
|
|
}
|
||
|
|
metrics_path = output_dir / "metrics.json"
|
||
|
|
metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
|
||
|
|
print_metrics("fused", fused_metrics)
|
||
|
|
print_metrics("frozen EgoSpeak", base_metrics)
|
||
|
|
print(f"\nSaved metrics: {metrics_path}")
|
||
|
|
print(f"Fused evaluator input: {output_dir}")
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|