From 64adbef13dfa0065f17e26f46fbf95f909866ae2 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Wed, 16 Sep 2026 17:04:46 +0800 Subject: [PATCH] Add selective EgoSpeak Omni stage-two finetuning --- PROJECT_CONTEXT.md | 13 + docs/README.md | 1 + docs/easycom_egospeak_omni_finetune.md | 188 +++++++++ docs/easycom_egospeak_omni_residual.md | 7 + egospeak_omni_finetune.py | 195 ++++++++++ test_egospeak_omni_finetune.py | 374 ++++++++++++++++++ train_egospeak_omni_finetune.py | 520 +++++++++++++++++++++++++ 7 files changed, 1298 insertions(+) create mode 100644 docs/easycom_egospeak_omni_finetune.md create mode 100644 egospeak_omni_finetune.py create mode 100644 test_egospeak_omni_finetune.py create mode 100755 train_egospeak_omni_finetune.py diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index 6c2a7a9..bce567f 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -154,6 +154,19 @@ result. states remains a future extraction ablation. See the [`runbook`](docs/easycom_egospeak_omni_residual.md). +### Unmeasured dual-stream continuation + +Stage-two `egospeak-omni-selective-finetune` is implemented but unmeasured. +It strictly continues the trained stage-one checkpoint, selectively adapts +MiniROADA's GRU/forecast heads with a lower learning rate, and leaves original +feature extractors, input projection and Omni frozen. It exports immutable +stage-one/original references and automatic threshold-free raw endpoint AP. +The user requested this exploratory continuation while semantic annotations +are audited; P0 remains unfinished. Four additional epochs means eight total +adaptation epochs, requiring the fusion-only continuation control for a +budget-matched unfreezing comparison. See the +[`stage-two runbook`](docs/easycom_egospeak_omni_finetune.md). + ### Measured oracle-information screen - A seed-0, four-epoch Session-12 development screen supplied transcript-derived diff --git a/docs/README.md b/docs/README.md index 6a14b52..e9d38dc 100644 --- a/docs/README.md +++ b/docs/README.md @@ -29,6 +29,7 @@ tests. | [`easycom_causal_person_adapter_v2.md`](easycom_causal_person_adapter_v2.md) | Dynamic-gated, zero-output causal person-track residual and matched L20 screen | | [`easycom_egospeak_style_person_gru.md`](easycom_egospeak_style_person_gru.md) | Controlled V2 temporal-stem ablation using EgoSpeak MiniROADA's layer order | | [`easycom_egospeak_omni_residual.md`](easycom_egospeak_omni_residual.md) | Stage-one frozen released MiniROADA plus zero-initialized tail-aware Omni residual fusion | +| [`easycom_egospeak_omni_finetune.md`](easycom_egospeak_omni_finetune.md) | Stage-two selective MiniROADA GRU/head continuation with immutable stage-one/base references and raw endpoint AP | | [`easycom_highrate_person_av.md`](easycom_highrate_person_av.md) | Frozen 25 Hz causal ASD person evidence fused with 5 Hz global Omni context | | [`easycom_muvap_stage3_baseline.md`](easycom_muvap_stage3_baseline.md) | MuVAP-only frozen VAP+ASD Stage-3 paper reconstruction on EgoSpeak/EasyCom targets | | [`easycom_oracle_information.md`](easycom_oracle_information.md) | Causal oracle per-person speaking-history upper bound and Session-12 screening commands | diff --git a/docs/easycom_egospeak_omni_finetune.md b/docs/easycom_egospeak_omni_finetune.md new file mode 100644 index 0000000..5f1c25e --- /dev/null +++ b/docs/easycom_egospeak_omni_finetune.md @@ -0,0 +1,188 @@ +# Stage two: selective MiniROADA + Omni residual fine-tuning + +## Decision and claim boundary + +The user requested this continuation after the +[`stage-one screen`](../results/easycom_egospeak_omni_residual_2026-09-16/README.md): +causal-valid dense mean mAP improved by +1.521 points, raw onset endpoint AP +by +0.600 and offset endpoint AP by +2.580. These are single-seed reported +results, not statistically significant semantic floor-control evidence. +Provisional ENTRY/YIELD evaluation is paused pending an annotation audit; +historical provisional results remain archived. P0 is still unfinished. + +This stage means selective MiniROADA temporal/forecast adaptation, **not** +Omni LoRA, a person-memory adapter, new feature extraction or a semantic head. +It is an exploratory implementation of the earlier selective-unfreezing route. +Do not use Sessions 1--3 to choose learning rates, windows, loss weights, +unfreezing policy or the stopping epoch. A comparison observed on the exposed +test split remains exploratory; choose any later settings on training-side +development data and report the test exposure honestly. + +## What changes + +Load `checkpoint_last.pt` from stage one strictly. Preserve its learned +nonzero residual and verify exact initial recovery of its fused/base logits +on a training example. The original released checkpoint is not needed again: +the first-stage checkpoint already contains its full frozen base parameters. + +| Component | Stage-two default | +|---|---| +| Released RGB and wav2vec2 feature extractors | Frozen; reuse existing arrays | +| MiniROADA RGB+audio input projection (`layer1`) | Frozen, dropout remains off | +| MiniROADA GRU | Trainable, learning rate 1e-5 | +| Anticipation projection and shared classification head | Trainable, learning rate 1e-5 | +| Ego/Omni projections, tail embedding and fusion residual | Continue training, learning rate 1e-4 | +| Qwen2.5-Omni backbone | Frozen; reuse exactly the first-stage sidecars | + +The GRU state contributes to both the original anticipation logits and the +feature-fusion residual. Unlike stage one, neither path is enclosed in +`no_grad` or detached when selected for training. The default trainable +MiniROADA modules contain 19,942,403 parameters, in addition to the existing +residual branch. The RGB/audio encoders themselves are not loaded or trained; +this is partial task-model fine-tuning with frozen feature extractors, not +raw-media end-to-end training. + +Loss, window/stride, grid, class order and ten horizons stay unchanged. +Window size and stride inherit the stage-one checkpoint and mismatched +overrides fail. Readout mode/layer fusion/dimension must also match; when the +stage-one sibling `metrics.json` is available, its full saved Omni extraction +signature is checked. A relocated checkpoint without that provenance prints +a warning, so the same original Omni root must be reused explicitly. + +Default optimizer: AdamW, weight decay 0.05, gradient clipping at norm 1.0. +These defaults are specified before the new run, not tuned on test results. +The first training batch logs separate fusion/base gradient norms. Final +evaluation occurs only after the fixed additional epoch budget, with no +test-side early stopping or best-checkpoint selection. + +## Session 4 smoke test + +From the repository root, reuse the same two feature roots as stage one: + +```bash +EGO_ROOT=third_party/EgoSpeak/MiniROAD/data/EasyCom +OMNI_ROOT=/path/to/the/same/first_stage/omni_features +STAGE1=results/easycom_egospeak_omni_residual_seed0/checkpoint_last.pt + +CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_finetune.py \ + --egospeak-feature-root "$EGO_ROOT" \ + --omni-feature-root "$OMNI_ROOT" \ + --stage-one-checkpoint "$STAGE1" \ + --train-sessions Session_4 \ + --eval-sessions Session_4 \ + --epochs 1 \ + --max-train-examples 32 \ + --batch-size 8 \ + --num-workers 0 \ + --output-dir results/easycom_egospeak_omni_finetune_smoke_2026-09-16 +``` + +Check exact initialization, nonzero fusion/base gradients, finite loss, +`[T,3,10]` predictions, `[T,10]` common valid masks and raw endpoint counts. +The printed cutoff is `(source_row+1)/5` s; input stops at that row, and +future labels are rows source+1 through source+10. Smoke predictions are not +the final model's calibration predictions. Start the formal run from the +original stage-one checkpoint, not the smoke checkpoint. + +## Fixed-budget continuation + +After the GPU smoke audit passes: + +```bash +CUDA_VISIBLE_DEVICES=0 python3 -u train_egospeak_omni_finetune.py \ + --egospeak-feature-root "$EGO_ROOT" \ + --omni-feature-root "$OMNI_ROOT" \ + --stage-one-checkpoint "$STAGE1" \ + --base-trainable gru-heads \ + --fusion-lr 1e-4 \ + --base-lr 1e-5 \ + --epochs 4 \ + --batch-size 64 \ + --num-workers 4 \ + --seed 0 \ + --output-dir results/easycom_egospeak_omni_finetune_seed0_2026-09-16 +``` + +Defaults use Sessions 4--12 for training and Sessions 1--3 for final +evaluation. Training or a first-stage checkpoint trained on Sessions 1--3 is +rejected. Float32 is the first reproducibility run; do not enable AMP until +its output/causality audit passes. Full-session GRU inference is the default; +`--eval-chunk-size N` carries recurrent state across chunks if memory requires. +Nonempty output directories are never overwritten. + +An additional four epochs means **4+4=8 total adaptation epochs**. It cannot +isolate unfreezing against the original four-epoch first stage. The required +matched-budget continuation control changes only these arguments: + +```bash + --base-trainable none \ + --output-dir results/easycom_egospeak_omni_fusion_continue_seed0_2026-09-16 +``` + +Keep the same stage-one checkpoint, seed, training split, four additional +epochs, sampling, optimizer and fusion learning rate. `--base-trainable heads` +is a separate predefined ablation that freezes the GRU and adapts only the +forecast heads plus residual; do not select among policies on test scores. + +## Outputs and reading AP + +```text +RESULT/ + README.md + checkpoint_last.pt # full adapted model, config, parent hash, budget + metrics.json # adapted dense + raw endpoint AP + raw_endpoint_metrics.csv # three methods, ten horizons, AP and counts + pred/, valid/, raw_eval/ # second-stage predictions + stage_one_reference/ + metrics.json, pred/, valid/, raw_eval/ + frozen_egospeak_reference/ + metrics.json, pred/, valid/, raw_eval/ +``` + +Both references are immutable copies of the first-stage checkpoint: one +includes its learned fusion, the other uses its original frozen MiniROADA +path. The adapted base branch is never mislabeled as frozen EgoSpeak. +Every method uses the same valid source rows. Reference directories are +directly compatible with the existing evaluator. A loader for the new full +model is `egospeak_omni_finetune.load_stage_two_checkpoint`; stage-one and +stage-two checkpoint formats are intentionally distinct. + +Dense metrics retain both full-timeline and causal-valid protocols. Raw onset +uses current wearer OFF, future ON versus future OFF with score p_wearer; +offset uses current wearer ON, future OFF versus future ON with score +1-p_wearer. AP pools eligible pairs over evaluated sessions separately for +each horizon; the headline averages ten APs. This is endpoint anticipation, +not timestamp detection or semantic ENTRY/YIELD. No Session 4 threshold +calibration or annotation file is needed for these APs. An undefined horizon +is preserved as null, not silently removed from the mean. + +The CSV/JSON store AP **fractions**, while console output uses percentages. +To print the raw results: + +```bash +python3 - <<'PY' +import pandas as pd +p = "results/easycom_egospeak_omni_finetune_seed0_2026-09-16/raw_endpoint_metrics.csv" +x = pd.read_csv(p) +cols = ["onset_endpoint_ap", "offset_endpoint_ap"] +x[cols] *= 100 +print(x[["run", "nominal_seconds", *cols]].round(3).to_string(index=False)) +print(x.groupby("run")[cols].mean().round(3).to_string()) +PY +``` + +If any horizon is undefined, use the JSON headline (which stays null) instead +of pandas' default mean that skips NaN. No semantic opportunity evaluation is +run by this launcher. + +## Local validation + +```bash +python3 -m unittest test_egospeak_omni_finetune test_egospeak_omni_residual -v +``` + +Tests cover trained residual preservation, selected-module gradients, +reference immutability, policy masks, causal perturbation, carried-state +chunk equivalence, checkpoint failure/success paths, raw AP and a tiny CPU +training/export/reload smoke. They do not substitute for the real Session 4 +GPU smoke or establish second-stage performance; that run remains unmeasured. diff --git a/docs/easycom_egospeak_omni_residual.md b/docs/easycom_egospeak_omni_residual.md index 82dc846..c3fdd35 100644 --- a/docs/easycom_egospeak_omni_residual.md +++ b/docs/easycom_egospeak_omni_residual.md @@ -154,6 +154,13 @@ so it remains readable by `summarize_local_baselines.py`. ## Interpretation and next gate +The first seed-0 dense/raw screen is now +[`recorded`](../results/easycom_egospeak_omni_residual_2026-09-16/README.md). +At the user's request, provisional semantic evaluation is paused for an +annotation audit. An exploratory selective GRU/head continuation is implemented +in the [`stage-two runbook`](easycom_egospeak_omni_finetune.md); this does not +imply that the semantic P0 gate below has passed. + This first stage uses only the released dense future-activity objective. It tests whether Omni adds useful context without discarding the task-specific EgoSpeak path. It does not train semantic ENTRY/YIELD heads, and Sessions diff --git a/egospeak_omni_finetune.py b/egospeak_omni_finetune.py new file mode 100644 index 0000000..7685736 --- /dev/null +++ b/egospeak_omni_finetune.py @@ -0,0 +1,195 @@ +"""Stage-two selective MiniROADA adaptation from a trained Omni residual. + +The released input features and Omni states remain frozen. This module only +enables gradients in selected MiniROADA temporal/forecast modules and the +existing residual. Stage-one checkpoints are loaded strictly, without resetting +the learned residual or requiring the original checkpoint to be available. +""" + +from __future__ import annotations + +from copy import deepcopy +from pathlib import Path +from typing import Any, Literal + +import torch + +from easycom_head import NUM_HORIZONS, OFFICIAL_TRAIN_SESSIONS +from egospeak_omni_residual import ( + EgoSpeakOmniResidual, + MiniROADABase, + load_torch_payload, + normalize_session_name, +) + +BaseTrainable = Literal["none", "heads", "gru-heads"] +STAGE_ONE_FORMAT = "egospeak-omni-zero-residual-v1" +STAGE_TWO_FORMAT = "egospeak-omni-selective-finetune-v1" + + +def validate_training_sessions(values: list[str]) -> None: + normalized = [normalize_session_name(value) for value in values] + if not normalized or len(normalized) != len(set(normalized)): + raise ValueError("Training sessions must be nonempty and unique") + if not set(normalized).issubset(OFFICIAL_TRAIN_SESSIONS): + raise ValueError("Training may only use EasyCom Sessions 4--12") + + +def load_stage_one_checkpoint( + path: Path, +) -> tuple[EgoSpeakOmniResidual, dict[str, Any]]: + path = path.expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(path) + payload = load_torch_payload(path) + if not isinstance(payload, dict) or payload.get("format") != STAGE_ONE_FORMAT: + raise ValueError( + "--stage-one-checkpoint must be a stage-one residual checkpoint" + ) + config = payload.get("model_config") + if not isinstance(config, dict): + raise TypeError("Stage-one checkpoint lacks model_config") + required = { + "omni_dim", + "fusion_dim", + "tail_embedding_dim", + "num_tail_bins", + "dropout", + } + if set(config) != required or config["num_tail_bins"] != NUM_HORIZONS: + raise ValueError("Incompatible stage-one residual model_config") + training = payload.get("training_config") + if not isinstance(training, dict) or int(training.get("epochs", 0)) <= 0: + raise ValueError("Stage-one checkpoint lacks a positive training epoch budget") + validate_training_sessions(payload.get("train_sessions", [])) + base_config = payload.get("base_model_config", {}) + if not isinstance(base_config, dict): + raise TypeError("Invalid base_model_config") + # Released v1 checkpoints omit base_model_config and use default MiniROADA. + model = EgoSpeakOmniResidual(MiniROADABase(**base_config), **config) + model.load_state_dict(payload["model_state_dict"], strict=True) + if any(not torch.isfinite(value).all() for value in model.state_dict().values()): + raise ValueError("Stage-one checkpoint contains nonfinite parameters") + model.requires_grad_(False) + model.eval() + return model, payload + + +class EgoSpeakOmniFinetune(EgoSpeakOmniResidual): + """Continuation with explicit base policy and intact stage-one logits.""" + + def __init__( + self, *args: Any, base_trainable: BaseTrainable = "gru-heads", **kwargs: Any + ): + super().__init__(*args, **kwargs) + if base_trainable not in ("none", "heads", "gru-heads"): + raise ValueError(f"Invalid base-trainable policy: {base_trainable}") + self.base_trainable = base_trainable + self.requires_grad_(True) + self.base.requires_grad_(False) + if base_trainable in ("heads", "gru-heads"): + self.base.anticipation_layer.requires_grad_(True) + self.base.f_classification.requires_grad_(True) + if base_trainable == "gru-heads": + self.base.gru.requires_grad_(True) + self.train(False) + + @classmethod + def from_stage_one( + cls, + reference: EgoSpeakOmniResidual, + config: dict[str, Any], + *, + base_trainable: BaseTrainable = "gru-heads", + ) -> EgoSpeakOmniFinetune: + model = cls(deepcopy(reference.base), **config, base_trainable=base_trainable) + # Copy rather than share weights with the immutable first-stage reference. + model.load_state_dict(reference.state_dict(), strict=True) + return model + + def train(self, mode: bool = True) -> EgoSpeakOmniFinetune: + super().train(mode) + # Keep frozen layer1's dropout off. Enabling GRU/head train mode does + # not turn the original frozen RGB+audio projection back on. + if self.base_trainable in ("heads", "gru-heads"): + self.base.anticipation_layer.train(mode) + self.base.f_classification.train(mode) + if self.base_trainable == "gru-heads": + self.base.gru.train(mode) + return self + + def optimizer_groups( + self, *, fusion_lr: float, base_lr: float + ) -> list[dict[str, Any]]: + if fusion_lr <= 0 or base_lr <= 0: + raise ValueError("Both learning rates must be positive") + fusion = [ + p + for name, p in self.named_parameters() + if p.requires_grad and not name.startswith("base.") + ] + base = [ + p + for name, p in self.named_parameters() + if p.requires_grad and name.startswith("base.") + ] + groups = [{"name": "fusion", "params": fusion, "lr": fusion_lr}] + if base: + groups.append({"name": "miniroada", "params": base, "lr": base_lr}) + return groups + + def forward_last( + self, + rgb: torch.Tensor, + audio: torch.Tensor, + omni: torch.Tensor, + tail_bin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + hidden, _ = self.base.encode(rgb, audio) + ego_hidden = hidden[:, -1, :] + base_logits = self.base.anticipation_logits(ego_hidden) + # No no_grad/detach: gradients must reach the selected MiniROADA GRU. + residual = self._residual_logits(ego_hidden, omni, tail_bin) + return base_logits + residual, base_logits + + def forward_sequence_chunk( + self, + rgb: torch.Tensor, + audio: torch.Tensor, + omni: torch.Tensor, + tail_bin: torch.Tensor, + recurrent_state: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + hidden, next_state = self.base.encode(rgb, audio, recurrent_state) + base_logits = self.base.anticipation_logits(hidden) + residual = self._residual_logits(hidden, omni, tail_bin) + return base_logits + residual, base_logits, next_state + + +def base_model_config(base: MiniROADABase) -> dict[str, Any]: + return { + "rgb_dim": base.rgb_dim, + "audio_dim": base.audio_dim, + "embedding_dim": base.embedding_dim, + "hidden_dim": base.hidden_dim, + "num_layers": base.num_layers, + "num_horizons": base.num_horizons, + "num_classes": base.num_classes, + "dropout": base.layer1[-1].p, + } + + +def load_stage_two_checkpoint( + path: Path, +) -> tuple[EgoSpeakOmniFinetune, dict[str, Any]]: + payload = load_torch_payload(path) + if not isinstance(payload, dict) or payload.get("format") != STAGE_TWO_FORMAT: + raise ValueError("Not a stage-two fine-tuning checkpoint") + model = EgoSpeakOmniFinetune( + MiniROADABase(**payload["base_model_config"]), + **payload["model_config"], + base_trainable=payload["base_trainable"], + ) + model.load_state_dict(payload["model_state_dict"], strict=True) + model.eval() + return model, payload diff --git a/test_egospeak_omni_finetune.py b/test_egospeak_omni_finetune.py new file mode 100644 index 0000000..72d0b7b --- /dev/null +++ b/test_egospeak_omni_finetune.py @@ -0,0 +1,374 @@ +import csv +import io +import json +import tempfile +import unittest +from contextlib import redirect_stdout +from pathlib import Path +from unittest.mock import patch + +import numpy as np +import torch + +from egospeak_omni_finetune import ( + STAGE_ONE_FORMAT, + STAGE_TWO_FORMAT, + EgoSpeakOmniFinetune, + base_model_config, + load_stage_one_checkpoint, + load_stage_two_checkpoint, + validate_training_sessions, +) +from egospeak_omni_residual import ( + EgoSpeakOmniResidual, + MiniROADABase, + load_easycom_session_arrays, + miniroad_normalized_loss, +) +from train_egospeak_omni_finetune import ( + main, + parse_args, + prepare_output, + raw_endpoint_metrics, + validate_checkpoint_contract, +) + + +def reference_model(): + torch.manual_seed(7) + base = MiniROADABase( + rgb_dim=4, audio_dim=6, embedding_dim=8, hidden_dim=5, dropout=0.4 + ) + config = { + "omni_dim": 7, + "fusion_dim": 4, + "tail_embedding_dim": 3, + "num_tail_bins": 10, + "dropout": 0.3, + } + model = EgoSpeakOmniResidual(base, **config) + # Simulate a trained, nonzero first-stage residual that must not be reset. + with torch.no_grad(): + model.residual_output.weight.normal_(std=0.1) + model.residual_output.bias.normal_(std=0.1) + model.requires_grad_(False) + model.eval() + return model, config + + +def stage_one_payload(reference, config): + return { + "format": STAGE_ONE_FORMAT, + "model_config": config, + "base_model_config": base_model_config(reference.base), + "model_state_dict": reference.state_dict(), + "train_sessions": ["Session_4"], + "training_config": { + "epochs": 4, + "window_size": 8, + "stride": 1, + "expected_omni_mode": "open-stream-replay", + "omni_layer_fusion": "mean", + }, + } + + +class EgoSpeakOmniFinetuneTest(unittest.TestCase): + def setUp(self): + # Small synthetic tests should not spin up many CPU BLAS workers. + self.previous_threads = torch.get_num_threads() + torch.set_num_threads(1) + + def tearDown(self): + torch.set_num_threads(self.previous_threads) + + def test_nonzero_stage_one_recovery_gradients_and_reference_immutability(self): + reference, config = reference_model() + model = EgoSpeakOmniFinetune.from_stage_one(reference, config) + inputs = ( + torch.randn(2, 9, 4), + torch.randn(2, 9, 6), + torch.randn(2, 7), + torch.tensor([0, 9]), + ) + expected, _ = reference.forward_last(*inputs) + actual, _ = model.forward_last(*inputs) + torch.testing.assert_close(actual, expected, rtol=0, atol=0) + self.assertGreater(float(reference.residual_output.weight.abs().sum()), 0) + original = {k: v.clone() for k, v in reference.state_dict().items()} + model.train() + self.assertFalse(model.base.layer1.training) + self.assertTrue(model.base.gru.training) + self.assertTrue(model.base.anticipation_layer.training) + logits, _ = model.forward_last(*inputs) + targets = torch.zeros_like(logits) + targets[..., 1] = 1 + loss = miniroad_normalized_loss(logits, targets) + loss.backward() + for module in ( + model.base.gru, + model.base.anticipation_layer, + model.base.f_classification, + model.residual_output, + model.omni_projection, + ): + norm = sum( + float(p.grad.abs().sum()) + for p in module.parameters() + if p.grad is not None + ) + self.assertGreater(norm, 0) + self.assertTrue(all(p.grad is None for p in model.base.layer1.parameters())) + self.assertTrue(all(p.grad is None for p in reference.parameters())) + groups = model.optimizer_groups(fusion_lr=1e-3, base_lr=1e-4) + ids = [id(p) for group in groups for p in group["params"]] + self.assertEqual(len(ids), len(set(ids))) + self.assertEqual(set(ids), {id(p) for p in model.trainable_parameters()}) + before = model.base.gru.weight_ih_l0.detach().clone() + torch.optim.AdamW(groups).step() + self.assertGreater( + float((before - model.base.gru.weight_ih_l0).abs().max().detach()), 0 + ) + for key, value in reference.state_dict().items(): + torch.testing.assert_close(value, original[key], rtol=0, atol=0) + + def test_freezing_policies(self): + reference, config = reference_model() + for policy in ("none", "heads", "gru-heads"): + with self.subTest(policy=policy): + model = EgoSpeakOmniFinetune.from_stage_one( + reference, config, base_trainable=policy + ) + model.train() + for name, parameter in model.base.named_parameters(): + expected = (name.startswith("gru.") and policy == "gru-heads") or ( + name.startswith(("anticipation_layer.", "f_classification.")) + and policy != "none" + ) + self.assertEqual(parameter.requires_grad, expected, name) + self.assertFalse(model.base.layer1[-1].training) + self.assertEqual( + len(model.optimizer_groups(fusion_lr=1e-4, base_lr=1e-5)), + 1 if policy == "none" else 2, + ) + with self.assertRaises(ValueError): + EgoSpeakOmniFinetune.from_stage_one(reference, config, base_trainable="all") + + def test_continuation_inherits_and_locks_window_and_readout(self): + reference, config = reference_model() + payload = stage_one_payload(reference, config) + required = [ + "--egospeak-feature-root", + "/tmp/ego", + "--omni-feature-root", + "/tmp/omni", + "--stage-one-checkpoint", + "/tmp/stage1.pt", + "--output-dir", + "/tmp/stage2", + ] + args = parse_args(required) + validate_checkpoint_contract(args, payload) + self.assertEqual((args.window_size, args.stride), (8, 1)) + for options in ( + ["--window-size", "16"], + ["--stride", "4"], + ["--omni-layer-fusion", "last"], + ["--expected-omni-mode", "any"], + ): + with self.subTest(options=options), self.assertRaises(ValueError): + validate_checkpoint_contract(parse_args(required + options), payload) + + def test_sequence_is_causal_and_chunked_recurrent_inference_matches(self): + reference, config = reference_model() + model = EgoSpeakOmniFinetune.from_stage_one(reference, config) + rgb, audio = torch.randn(1, 15, 4), torch.randn(1, 15, 6) + omni, tail = torch.randn(1, 15, 7), torch.arange(15)[None] % 10 + with torch.no_grad(): + full, _, _ = model.forward_sequence_chunk(rgb, audio, omni, tail, None) + prefix, _, state = model.forward_sequence_chunk( + rgb[:, :6], audio[:, :6], omni[:, :6], tail[:, :6], None + ) + suffix, _, _ = model.forward_sequence_chunk( + rgb[:, 6:], audio[:, 6:], omni[:, 6:], tail[:, 6:], state + ) + torch.testing.assert_close( + full, torch.cat((prefix, suffix), dim=1), rtol=1e-5, atol=1e-6 + ) + rgb[:, 6:] += 100 + audio[:, 6:] -= 100 + omni[:, 6:] += 100 + changed, _, _ = model.forward_sequence_chunk(rgb, audio, omni, tail, None) + torch.testing.assert_close(full[:, :6], changed[:, :6], rtol=0, atol=0) + + def test_checkpoint_loaders_reject_wrong_formats_test_training_and_nonfinite(self): + reference, config = reference_model() + with tempfile.TemporaryDirectory() as directory: + path = Path(directory) / "checkpoint.pt" + payload = stage_one_payload(reference, config) + torch.save(payload, path) + loaded, _ = load_stage_one_checkpoint(path) + torch.testing.assert_close( + loaded.residual_output.weight, reference.residual_output.weight + ) + for change in ( + {"format": STAGE_TWO_FORMAT}, + {"train_sessions": ["Session_1"]}, + {"model_config": {**config, "num_tail_bins": 9}}, + ): + torch.save({**payload, **change}, path) + with self.assertRaises(ValueError): + load_stage_one_checkpoint(path) + with torch.no_grad(): + payload["model_state_dict"]["residual_output.bias"][0] = float("nan") + torch.save(payload, path) + with self.assertRaisesRegex(ValueError, "nonfinite"): + load_stage_one_checkpoint(path) + with self.assertRaises(FileNotFoundError): + load_stage_one_checkpoint(Path("/does-not-exist/checkpoint.pt")) + for sessions in (["Session_1"], ["Session_4", "Session_4"], []): + with self.assertRaises(ValueError): + validate_training_sessions(sessions) + + def test_raw_ap_uses_endpoint_labels_and_has_no_calibration_requirement(self): + targets = np.zeros((14, 3), dtype=np.float32) + targets[:, 0] = 1 + wearer = np.array([0, 0, 1, 1] * 3 + [0, 0], dtype=bool) + targets[wearer, 0] = 0 + targets[wearer, 1] = 1 + predictions = np.zeros((14, 3, 10), dtype=np.float32) + predictions[:, 1, :] = targets[:, 1, None] + valid = np.ones((14, 10), dtype=bool) + metrics = raw_endpoint_metrics( + {"Session_1": predictions}, {"Session_1": valid}, {"Session_1": targets} + ) + self.assertEqual(metrics["horizons"][0]["onset"]["average_precision"], 1) + self.assertEqual(metrics["horizons"][0]["offset"]["average_precision"], 1) + # Four-step horizons have no positives, so headline stays undefined. + self.assertIsNone(metrics["mean_onset_ap"]) + + def test_cpu_end_to_end_smoke_export_reload_and_fixed_reference_predictions(self): + reference, config = reference_model() + with tempfile.TemporaryDirectory() as directory: + root = Path(directory) + checkpoint = root / "stage1" / "checkpoint_last.pt" + checkpoint.parent.mkdir() + torch.save(stage_one_payload(reference, config), checkpoint) + ego, omni = root / "ego", root / "omni" + for sub in ("rgb_kinetics_resnet50", "audio_wav2vec2", "target_perframe"): + (ego / sub).mkdir(parents=True) + omni.mkdir() + steps = 36 + rng = np.random.default_rng(3) + np.save( + ego / "rgb_kinetics_resnet50" / "Session_4.npy", + rng.standard_normal((steps, 4)).astype(np.float32), + ) + np.save( + ego / "audio_wav2vec2" / "Session_4.npy", + rng.standard_normal((steps, 6)).astype(np.float32), + ) + target = np.eye(3, dtype=np.float32)[np.arange(steps) % 3] + np.save(ego / "target_perframe" / "Session_4.npy", target) + torch.save( + { + "mode": "open-stream-replay", + "states": torch.randn(steps, 2, 7), + "times": (torch.arange(steps, dtype=torch.float64) + 1) / 5, + "label_indices": torch.arange(steps), + "start": 0, + "config": {"step_seconds": 0.2, "native_block_seconds": 2.0}, + }, + omni / "Session_4.pt", + ) + session = load_easycom_session_arrays( + ego, omni, "Session_4", rgb_dim=4, audio_dim=6, omni_layer_fusion="mean" + ) + output = root / "stage2" + arguments = [ + "--egospeak-feature-root", + str(ego), + "--omni-feature-root", + str(omni), + "--stage-one-checkpoint", + str(checkpoint), + "--output-dir", + str(output), + "--train-sessions", + "Session_4", + "--eval-sessions", + "Session_4", + "--device", + "cpu", + "--epochs", + "1", + "--batch-size", + "2", + "--num-workers", + "0", + "--max-train-examples", + "4", + "--no-progress", + ] + with ( + patch( + "train_egospeak_omni_finetune.load_sessions", return_value=[session] + ), + redirect_stdout(io.StringIO()), + ): + main(arguments) + loaded, payload = load_stage_two_checkpoint(output / "checkpoint_last.pt") + self.assertEqual(payload["budgets"]["total_adaptation_epochs"], 5) + self.assertGreater( + payload["history"][0]["first_batch_gradient_norms"]["miniroada"], 0 + ) + self.assertEqual(loaded.base_trainable, "gru-heads") + self.assertTrue((output / "README.md").is_file()) + for method in ( + output, + output / "stage_one_reference", + output / "frozen_egospeak_reference", + ): + pred = np.load(method / "pred" / "Session_4.npy") + valid = np.load(method / "valid" / "Session_4.npy") + self.assertEqual(pred.shape, (steps, 3, 10)) + self.assertEqual(valid.shape, (steps, 10)) + self.assertEqual(valid.sum(), (steps - 10) * 10) + self.assertTrue(np.isfinite(pred).all()) + np.testing.assert_array_equal( + valid, np.load(output / "valid" / "Session_4.npy") + ) + with torch.no_grad(): + expected_stage1, expected_original, _ = ( + reference.forward_sequence_chunk( + torch.from_numpy(np.asarray(session.rgb).copy())[ + None, : steps - 10 + ], + torch.from_numpy(np.asarray(session.audio).copy())[ + None, : steps - 10 + ], + session.omni[None, : steps - 10], + session.tail_bins[None, : steps - 10], + None, + ) + ) + for name, expected in ( + ("stage_one_reference", expected_stage1), + ("frozen_egospeak_reference", expected_original), + ): + np.testing.assert_allclose( + np.load(output / name / "raw_eval" / "Session_4.npy"), + expected.softmax(-1)[0].numpy(), + rtol=0, + atol=0, + ) + metrics = json.loads((output / "metrics.json").read_text()) + self.assertIn("raw_endpoint_anticipation", metrics) + with (output / "raw_endpoint_metrics.csv").open() as handle: + self.assertEqual(len(list(csv.DictReader(handle))), 30) + with self.assertRaises(FileExistsError): + prepare_output(output) + + +if __name__ == "__main__": + unittest.main() diff --git a/train_egospeak_omni_finetune.py b/train_egospeak_omni_finetune.py new file mode 100755 index 0000000..3f96acf --- /dev/null +++ b/train_egospeak_omni_finetune.py @@ -0,0 +1,520 @@ +#!/usr/bin/env python3 +"""Continue stage-one fusion with selective MiniROADA GRU/head fine-tuning. + +Frozen RGB/audio features and causal Omni states are reused. Final evaluation +exports stage-two, immutable stage-one and original EgoSpeak predictions under +identical masks, plus dense and threshold-free raw endpoint AP. No semantic +annotations, dev checkpoint selection or test-side tuning are used. +""" + +from __future__ import annotations + +import argparse +import csv +import json +import math +from pathlib import Path +from typing import Any + +import numpy as np +import torch +from tqdm.auto import tqdm + +from easycom_head import ( + CLASS_NAMES, + NUM_HORIZONS, + OFFICIAL_TEST_SESSIONS, + OFFICIAL_TRAIN_SESSIONS, +) +from egospeak_omni_finetune import ( + STAGE_TWO_FORMAT, + EgoSpeakOmniFinetune, + base_model_config, + load_stage_one_checkpoint, + validate_training_sessions, +) +from egospeak_omni_residual import ( + EgoSpeakOmniResidual, + EgoSpeakOmniWindowDataset, + align_source_predictions, + miniroad_normalized_loss, + sha256_file, +) +from evaluate_floor_state_changes import evaluate_horizon +from train_egospeak_omni_residual import ( + autocast_context, + evaluate_predictions, + load_sessions, + make_loader, + normalize_sessions, + omni_contract_signature, + predict_session, + print_metrics, + resolve_device, + set_seed, +) + + +def parse_args(argv: list[str] | None = None) -> argparse.Namespace: + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--egospeak-feature-root", type=Path, required=True) + parser.add_argument("--omni-feature-root", type=Path, required=True) + parser.add_argument("--stage-one-checkpoint", type=Path, required=True) + 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") + parser.add_argument( + "--omni-layer-fusion", choices=("single", "last", "mean"), default="mean" + ) + parser.add_argument( + "--base-trainable", + choices=("none", "heads", "gru-heads"), + default="gru-heads", + help="none is a matched extra-epoch fusion-only control; layer1 always stays frozen.", + ) + parser.add_argument( + "--epochs", type=int, default=4, help="Additional stage-two epochs." + ) + parser.add_argument("--batch-size", type=int, default=64) + parser.add_argument("--num-workers", type=int, default=4) + parser.add_argument( + "--window-size", type=int, default=None, help="Inherit stage-one window." + ) + parser.add_argument( + "--stride", type=int, default=None, help="Inherit stage-one sampling stride." + ) + parser.add_argument("--fusion-lr", type=float, default=1e-4) + parser.add_argument("--base-lr", type=float, default=1e-5) + parser.add_argument("--weight-decay", type=float, default=0.05) + parser.add_argument("--grad-clip", type=float, default=1.0) + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--device", default="cuda") + parser.add_argument("--amp", action="store_true") + parser.add_argument("--eval-chunk-size", type=int, default=0) + parser.add_argument("--max-train-examples", type=int, default=None) + parser.add_argument("--no-progress", action="store_true") + return parser.parse_args(argv) + + +def validate_checkpoint_contract( + args: argparse.Namespace, payload: dict[str, Any] +) -> None: + training = payload["training_config"] + for key in ("omni_layer_fusion", "expected_omni_mode"): + if key not in training or getattr(args, key) != training[key]: + raise ValueError(f"Stage-two {key} must match stage-one extraction/readout") + if args.window_size is None: + args.window_size = int(training["window_size"]) + if args.stride is None: + args.stride = int(training["stride"]) + if args.window_size != training["window_size"] or args.stride != training["stride"]: + raise ValueError( + "Keep stage-one window_size/stride fixed for this continuation" + ) + + +@torch.inference_mode() +def verify_stage_one_recovery( + model: EgoSpeakOmniFinetune, + reference: EgoSpeakOmniResidual, + dataset: EgoSpeakOmniWindowDataset, + device: torch.device, +) -> float: + model.eval() + reference.eval() + rgb, audio, omni, tail, _, _, source = dataset[0] + inputs = ( + rgb[None].to(device), + audio[None].to(device), + omni[None].to(device), + tail[None].to(device), + ) + actual, actual_base = model.forward_last(*inputs) + expected, expected_base = reference.forward_last(*inputs) + maximum = max( + float((actual - expected).abs().max()), + float((actual_base - expected_base).abs().max()), + ) + if maximum != 0.0: + raise RuntimeError( + f"Stage-two initialization changed stage-one logits: {maximum}" + ) + print("Initialization audit: stage-two logits == trained stage-one logits exactly") + print( + f"Causal sample: source_row={int(source)}, cutoff={(int(source) + 1) / 5:.2f}s; " + "input ends at source row, labels are source+1,...,source+10" + ) + return maximum + + +def train_epoch(model, loader, optimizer, scaler, device, args) -> dict[str, Any]: + model.train() + total_loss = 0.0 + num_examples = 0 + first_batch_gradient_norms = {} + iterator = tqdm( + loader, desc="train stage2", disable=args.no_progress, dynamic_ncols=True + ) + for rgb, audio, omni, tail, targets, _, _ in iterator: + rgb = rgb.to(device=device, dtype=torch.float32) + audio = audio.to(device=device, dtype=torch.float32) + omni = omni.to(device=device, dtype=torch.float32) + tail = tail.to(device) + targets = targets.to(device=device, dtype=torch.float32) + optimizer.zero_grad(set_to_none=True) + with autocast_context(device, args.amp): + logits, _ = model.forward_last(rgb, audio, omni, tail) + loss = miniroad_normalized_loss(logits.float(), targets) + if not torch.isfinite(loss): + raise RuntimeError("Nonfinite stage-two training loss") + scaler.scale(loss).backward() + scaler.unscale_(optimizer) + if not num_examples: + for group in optimizer.param_groups: + norms = [ + p.grad.detach().norm() + for p in group["params"] + if p.grad is not None + ] + norm = float(torch.stack(norms).norm()) if norms else 0.0 + first_batch_gradient_norms[group["name"]] = norm + print( + f"First-batch gradient norms (before clipping): {first_batch_gradient_norms}" + ) + torch.nn.utils.clip_grad_norm_( + model.trainable_parameters(), + args.grad_clip, + error_if_nonfinite=True, + ) + scaler.step(optimizer) + scaler.update() + num_examples += len(rgb) + total_loss += float(loss.detach()) * len(rgb) + iterator.set_postfix_str(f"loss={total_loss / num_examples:.5f}") + return { + "loss": total_loss / num_examples, + "num_examples": num_examples, + "first_batch_gradient_norms": first_batch_gradient_norms, + } + + +def raw_endpoint_metrics(predictions, valid, targets) -> dict[str, Any]: + names = list(predictions) + target = np.concatenate([targets[name] for name in names]) + prediction = np.concatenate([predictions[name] for name in names]) + mask = np.concatenate([valid[name] for name in names]) + session_ids = np.concatenate( + [np.full(len(targets[name]), i) for i, name in enumerate(names)] + ) + horizons = [] + for horizon in range(NUM_HORIZONS): + item = evaluate_horizon(target, prediction, mask, session_ids, horizon, []) + horizons.append( + { + "seconds": item["nominal_seconds"], + "onset": item["wearer_off_to_on_endpoint"], + "offset": item["wearer_on_to_off_endpoint"], + } + ) + means = {} + for family in ("onset", "offset"): + values = [h[family]["average_precision"] for h in horizons] + # Undefined horizons stay undefined rather than being silently dropped. + means[f"mean_{family}_ap"] = ( + None if any(v is None for v in values) else float(np.mean(values)) + ) + return { + "protocol": "causal_valid_endpoint_state_change_ap", + "horizons": horizons, + **means, + } + + +def prepare_output(path: Path) -> Path: + path = path.expanduser().resolve() + if path.exists() and (not path.is_dir() or any(path.iterdir())): + raise FileExistsError(f"Refusing to overwrite output: {path}") + path.mkdir(parents=True, exist_ok=True) + for method in ( + path, + path / "stage_one_reference", + path / "frozen_egospeak_reference", + ): + for name in ("raw_eval", "pred", "valid"): + (method / name).mkdir(parents=True) + return path + + +def main(argv: list[str] | None = None) -> None: + args = parse_args(argv) + if ( + args.epochs <= 0 + or args.batch_size <= 0 + or args.eval_chunk_size < 0 + or args.num_workers < 0 + or args.grad_clip <= 0 + or args.weight_decay < 0 + or not all( + math.isfinite(value) + for value in ( + args.fusion_lr, + args.base_lr, + args.grad_clip, + args.weight_decay, + ) + ) + ): + raise ValueError("Invalid stage-two epoch/batch/chunk/worker/gradient settings") + set_seed(args.seed) + device = resolve_device(args.device) + if args.amp and device.type != "cuda": + raise ValueError("--amp requires CUDA") + reference, parent = load_stage_one_checkpoint(args.stage_one_checkpoint) + validate_checkpoint_contract(args, parent) + train_names = normalize_sessions(args.train_sessions, OFFICIAL_TRAIN_SESSIONS) + eval_names = normalize_sessions(args.eval_sessions, OFFICIAL_TEST_SESSIONS) + validate_training_sessions(train_names) + train_sessions = load_sessions(args, train_names) + eval_sessions = load_sessions(args, eval_names) + sessions = (*train_sessions, *eval_sessions) + if {int(s.omni.shape[-1]) for s in sessions} != {reference.omni_dim}: + raise ValueError("Omni feature dimensions differ from stage-one checkpoint") + signatures = {omni_contract_signature(s) for s in sessions} + if len(signatures) != 1: + raise ValueError("Train/evaluation Omni extraction contracts differ") + # v1 saves readout/mode but not the entire extraction contract in its + # checkpoint. Verify sibling metrics provenance when it is available. + parent_metrics = ( + args.stage_one_checkpoint.expanduser().resolve().parent / "metrics.json" + ) + contract_verified = False + if parent_metrics.is_file(): + metadata = json.loads(parent_metrics.read_text()).get("omni_sources", {}) + for source in metadata.values(): + config = source["config"] + signature = ( + source["model_id"], + source["mode"], + config.get("memory_seconds"), + config.get("video_fps"), + config.get("readout_last_n_layers"), + config.get("readout_layer_fusion"), + tuple(source["state_shape_before_layer_fusion"][1:]), + ) + if signature not in signatures: + raise ValueError( + "Current Omni extraction differs from stage-one metrics provenance" + ) + contract_verified = bool(metadata) + if not contract_verified: + print( + "WARNING: stage-one full extraction provenance unavailable; reuse its exact Omni root." + ) + + model = EgoSpeakOmniFinetune.from_stage_one( + reference, + parent["model_config"], + base_trainable=args.base_trainable, + ).to(device) + reference = reference.to(device) + groups = model.optimizer_groups(fusion_lr=args.fusion_lr, base_lr=args.base_lr) + for group in groups: + print( + f"parameters {group['name']}: {sum(p.numel() for p in group['params']):,}; lr={group['lr']:g}" + ) + print(f"base_trainable={args.base_trainable}; layer1 frozen=True; Omni frozen=True") + print(f"train={train_names}; final evaluation={eval_names}") + dataset = EgoSpeakOmniWindowDataset( + train_sessions, + window_size=args.window_size, + stride=args.stride, + seed=args.seed, + max_examples=args.max_train_examples, + ) + difference = verify_stage_one_recovery(model, reference, dataset, device) + output = prepare_output(args.output_dir) + optimizer = torch.optim.AdamW(groups, weight_decay=args.weight_decay) + try: + scaler = torch.amp.GradScaler("cuda", enabled=args.amp) + except (AttributeError, TypeError): # Older PyTorch compatibility. + scaler = torch.cuda.amp.GradScaler(enabled=args.amp) + history = [] + for epoch in range(1, args.epochs + 1): + dataset.resample(epoch - 1) + result = train_epoch( + model, make_loader(dataset, args, epoch), optimizer, scaler, device, args + ) + result["epoch"] = epoch + history.append(result) + print( + f"stage2 epoch {epoch}/{args.epochs}: loss={result['loss']:.6f}; examples={result['num_examples']}" + ) + + training_config = vars(args).copy() + training_config = { + k: str(v) if isinstance(v, Path) else v for k, v in training_config.items() + } + budgets = { + "stage_one_epochs": int(parent["training_config"]["epochs"]), + "stage_two_epochs": args.epochs, + "total_adaptation_epochs": int(parent["training_config"]["epochs"]) + + args.epochs, + } + checkpoint = output / "checkpoint_last.pt" + provenance = { + "stage_one_checkpoint": str(args.stage_one_checkpoint.expanduser().resolve()), + "stage_one_checkpoint_sha256": sha256_file( + args.stage_one_checkpoint.expanduser().resolve() + ), + "original_base_checkpoint": parent.get("base_checkpoint"), + "original_base_checkpoint_sha256": parent.get("base_checkpoint_sha256"), + "stage_one_full_extraction_contract_verified": contract_verified, + } + (output / "README.md").write_text( + "# Selective EgoSpeak + Omni fine-tuning\n\n" + f"Train: {train_names}. Final evaluation: {eval_names}. Seed: {args.seed}.\n\n" + f"Checkpoint: `checkpoint_last.pt`. Base policy: `{args.base_trainable}`. " + "RGB/audio feature extractors, input projection and Omni remain frozen.\n\n" + f"Epoch budgets: {budgets}. This continuation is not budget-matched " + "against stage one; use the `--base-trainable none` continuation control.\n\n" + "`metrics.json` reports dense AP/mAP and threshold-free raw endpoint AP " + "as fractions. `raw_endpoint_metrics.csv` has per-horizon AP and counts. " + "Raw onset uses OFF->ON versus OFF->OFF, scored by wearer probability; " + "offset uses ON->OFF versus ON->ON, scored by 1-wearer probability.\n\n" + "Immutable reference predictions are in `stage_one_reference/` and " + "`frozen_egospeak_reference/`, each with pred/valid/raw_eval and metrics. " + "No provisional semantic annotations are used.\n\n" + f"Run arguments:\n```json\n{json.dumps(training_config, indent=2)}\n```\n" + ) + torch.save( + { + "format": STAGE_TWO_FORMAT, + "model_state_dict": model.state_dict(), + "base_model_config": base_model_config(model.base), + "model_config": parent["model_config"], + "base_trainable": args.base_trainable, + "training_config": training_config, + "train_sessions": train_names, + "eval_sessions": eval_names, + "stage_one_train_sessions": parent["train_sessions"], + "budgets": budgets, + "provenance": provenance, + "omni_sources": {item.name: item.omni_metadata for item in sessions}, + "history": history, + "initialization_max_abs_stage_one_logit_difference": difference, + }, + checkpoint, + ) + + directories = { + "stage_two": output, + "stage_one": output / "stage_one_reference", + "egospeak": output / "frozen_egospeak_reference", + } + predictions = {name: {} for name in directories} + valid = {} + targets = {} + for session in eval_sessions: + kwargs = { + "chunk_size": args.eval_chunk_size, + "amp": args.amp, + "show_progress": not args.no_progress, + } + fused, _ = predict_session(model, session, device, **kwargs) + first_stage, original = predict_session(reference, session, device, **kwargs) + for method, raw in ( + ("stage_two", fused), + ("stage_one", first_stage), + ("egospeak", original), + ): + aligned, mask = align_source_predictions(raw, session.num_steps) + if session.name in valid and not np.array_equal(valid[session.name], mask): + raise RuntimeError("Reference and adapted valid masks differ") + valid[session.name] = mask + predictions[method][session.name] = aligned + directory = directories[method] + np.save(directory / "raw_eval" / f"{session.name}.npy", raw) + np.save(directory / "pred" / f"{session.name}.npy", aligned) + np.save(directory / "valid" / f"{session.name}.npy", mask) + targets[session.name] = np.asarray(session.targets, dtype=np.float32) + + dense = { + name: evaluate_predictions(pred, valid, targets) + for name, pred in predictions.items() + } + endpoints = { + name: raw_endpoint_metrics(pred, valid, targets) + for name, pred in predictions.items() + } + for name, directory in directories.items(): + metrics = { + **dense[name], + "format": "egospeak-omni-stage-two-eval-v1", + "classes": list(CLASS_NAMES), + "sessions": eval_names, + "train_sessions": train_names, + "checkpoint": str(checkpoint) + if name == "stage_two" + else provenance["stage_one_checkpoint"], + "method": name, + "seed": args.seed, + "prediction_layout": "[target_time, class, future_horizon]", + "raw_prediction_layout": "[source_time, future_horizon, class]", + "raw_endpoint_anticipation": endpoints[name], + } + if name == "stage_two": + metrics.update( + provenance=provenance, + budgets=budgets, + base_trainable=args.base_trainable, + training_config=training_config, + omni_sources={item.name: item.omni_metadata for item in sessions}, + stage_one_reference=dense["stage_one"], + frozen_egospeak_baseline=dense["egospeak"], + semantic_evaluation="not run: provisional annotation audit pending", + ) + (directory / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n") + print_metrics(name, dense[name]) + csv_path = output / "raw_endpoint_metrics.csv" + with csv_path.open("w", newline="") as handle: + fields = [ + "run", + "nominal_seconds", + "onset_endpoint_ap", + "offset_endpoint_ap", + "onset_positive", + "onset_negative", + "offset_positive", + "offset_negative", + ] + writer = csv.DictWriter(handle, fieldnames=fields) + writer.writeheader() + for name, metrics in endpoints.items(): + for item in metrics["horizons"]: + writer.writerow( + { + "run": name, + "nominal_seconds": item["seconds"], + "onset_endpoint_ap": item["onset"]["average_precision"], + "offset_endpoint_ap": item["offset"]["average_precision"], + "onset_positive": item["onset"]["num_positive"], + "onset_negative": item["onset"]["num_negative"], + "offset_positive": item["offset"]["num_positive"], + "offset_negative": item["offset"]["num_negative"], + } + ) + print("\nRaw endpoint ten-horizon means (percent; threshold-free):") + for name, metrics in endpoints.items(): + print( + name, + { + k: None if metrics[k] is None else round(metrics[k] * 100, 3) + for k in ("mean_onset_ap", "mean_offset_ap") + }, + ) + print(f"Saved {checkpoint}; raw endpoint CSV (fractions)={csv_path}") + + +if __name__ == "__main__": + main()