Add inference-only evaluation for EgoSpeak Omni finetuning

2026-09-16 17:28:31 +08:00
parent 64adbef13d
commit 7344d0c469
3 changed files with 181 additions and 0 deletions
@@ -177,6 +177,26 @@ run by this launcher.

 ## Local validation

+### Evaluate a saved checkpoint without training again
+
+For a completed short continuation, use the separate inference-only CLI:
+
+```bash
+CUDA_VISIBLE_DEVICES=0 python3 -u infer_egospeak_omni_finetune.py \
+  --checkpoint results/easycom_egospeak_omni_finetune_low_lr_1ep_diagnostic_2026-09-16/checkpoint_last.pt \
+  --sessions Session_1 Session_2 Session_3 \
+  --output-dir results/easycom_egospeak_omni_finetune_low_lr_1ep_test_2026-09-16
+```
+
+Feature roots/readout settings inherit the checkpoint; optional
+`--egospeak-feature-root`/`--omni-feature-root` override relocated roots while
+saved extraction contracts are checked. Float32 inference, `pred/valid/raw_eval`,
+dense and threshold-free raw endpoint AP are exported. No optimizer is created
+and the checkpoint is unchanged. This repeated official-test evaluation is
+exploratory; do not choose subsequent parameters or epochs from its scores.
+
+### Tests
+
 ```bash
 python3 -m unittest test_egospeak_omni_finetune test_egospeak_omni_residual -v
 ```
@@ -0,0 +1,128 @@
+"""Evaluate an existing stage-two checkpoint without any optimization."""
+
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import numpy as np
+
+from easycom_head import OFFICIAL_TEST_SESSIONS
+from egospeak_omni_finetune import load_stage_two_checkpoint
+from egospeak_omni_residual import align_source_predictions, sha256_file
+from train_egospeak_omni_finetune import raw_endpoint_metrics
+from train_egospeak_omni_residual import (
+    evaluate_predictions,
+    load_sessions,
+    normalize_sessions,
+    omni_contract_signature,
+    predict_session,
+    print_metrics,
+    resolve_device,
+)
+
+
+def main(argv: list[str] | None = None) -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--checkpoint", type=Path, required=True)
+    parser.add_argument("--output-dir", type=Path, required=True)
+    parser.add_argument("--sessions", nargs="+", default=None)
+    parser.add_argument("--egospeak-feature-root", type=Path)
+    parser.add_argument("--omni-feature-root", type=Path)
+    parser.add_argument("--device", default="cuda")
+    parser.add_argument("--eval-chunk-size", type=int, default=0)
+    parser.add_argument("--no-progress", action="store_true")
+    args = parser.parse_args(argv)
+    if args.eval_chunk_size < 0:
+        raise ValueError("Evaluation chunk size must be nonnegative")
+    model, payload = load_stage_two_checkpoint(args.checkpoint)
+    config = payload["training_config"]
+    args.egospeak_feature_root = args.egospeak_feature_root or Path(
+        config["egospeak_feature_root"]
+    )
+    args.omni_feature_root = args.omni_feature_root or Path(config["omni_feature_root"])
+    args.omni_layer_fusion = config["omni_layer_fusion"]
+    args.expected_omni_mode = config["expected_omni_mode"]
+    names = normalize_sessions(args.sessions, OFFICIAL_TEST_SESSIONS)
+    sessions = load_sessions(args, names)
+    if {int(s.omni.shape[-1]) for s in sessions} != {model.omni_dim}:
+        raise ValueError("Inference Omni dimension differs from checkpoint")
+    known = payload.get("omni_sources", {})
+    for session in sessions:
+        signature = omni_contract_signature(session)
+        for source in known.values():
+            c = source["config"]
+            expected = (
+                source["model_id"],
+                source["mode"],
+                c.get("memory_seconds"),
+                c.get("video_fps"),
+                c.get("readout_last_n_layers"),
+                c.get("readout_layer_fusion"),
+                tuple(source["state_shape_before_layer_fusion"][1:]),
+            )
+            if signature != expected:
+                raise ValueError(
+                    "Inference Omni extraction differs from checkpoint provenance"
+                )
+    output = args.output_dir.expanduser().resolve()
+    if output.exists() and (not output.is_dir() or any(output.iterdir())):
+        raise FileExistsError(f"Refusing to overwrite output: {output}")
+    for name in ("pred", "valid", "raw_eval"):
+        (output / name).mkdir(parents=True)
+    device = resolve_device(args.device)
+    model.requires_grad_(False)
+    model.to(device).eval()
+    predictions, valid, targets = {}, {}, {}
+    for session in sessions:
+        raw, _ = predict_session(
+            model,
+            session,
+            device,
+            chunk_size=args.eval_chunk_size,
+            amp=False,
+            show_progress=not args.no_progress,
+        )
+        aligned, mask = align_source_predictions(raw, session.num_steps)
+        predictions[session.name], valid[session.name] = aligned, mask
+        targets[session.name] = np.asarray(session.targets, dtype=np.float32)
+        np.save(output / "raw_eval" / f"{session.name}.npy", raw)
+        np.save(output / "pred" / f"{session.name}.npy", aligned)
+        np.save(output / "valid" / f"{session.name}.npy", mask)
+    dense = evaluate_predictions(predictions, valid, targets)
+    endpoints = raw_endpoint_metrics(predictions, valid, targets)
+    metrics = {
+        **dense,
+        "raw_endpoint_anticipation": endpoints,
+        "sessions": names,
+        "train_sessions": payload["train_sessions"],
+        "seed": config["seed"],
+        "checkpoint": str(args.checkpoint.resolve()),
+        "checkpoint_sha256": sha256_file(args.checkpoint),
+        "inference_only": True,
+        "budgets": payload["budgets"],
+        "omni_sources": {s.name: s.omni_metadata for s in sessions},
+    }
+    (output / "metrics.json").write_text(json.dumps(metrics, indent=2) + "\n")
+    (output / "README.md").write_text(
+        "# Stage-two inference-only evaluation\n\n"
+        f"Checkpoint: {args.checkpoint}. Sessions: {names}. No optimization.\n\n"
+        "metrics.json stores dense/raw endpoint AP as fractions; console uses percent. "
+        "Repeated official-test evaluation is exploratory, not dev selection.\n"
+    )
+    print_metrics("stage_two inference only", dense)
+    print(
+        "Raw endpoint mean AP (percent):",
+        {
+            family: None
+            if endpoints[f"mean_{family}_ap"] is None
+            else endpoints[f"mean_{family}_ap"] * 100
+            for family in ("onset", "offset")
+        },
+    )
+    print(f"Saved {output / 'metrics.json'}; checkpoint unchanged")
+
+
+if __name__ == "__main__":
+    main()
@@ -24,7 +24,9 @@ from egospeak_omni_residual import (
    MiniROADABase,
    load_easycom_session_arrays,
    miniroad_normalized_loss,
+    sha256_file,
 )
+from infer_egospeak_omni_finetune import main as infer_main
 from train_egospeak_omni_finetune import (
    main,
    parse_args,
@@ -324,6 +326,37 @@ class EgoSpeakOmniFinetuneTest(unittest.TestCase):
            )
            self.assertEqual(loaded.base_trainable, "gru-heads")
            self.assertTrue((output / "README.md").is_file())
+            infer_output = root / "infer_only"
+            before_hash = sha256_file(output / "checkpoint_last.pt")
+            with (
+                patch(
+                    "infer_egospeak_omni_finetune.load_sessions", return_value=[session]
+                ),
+                redirect_stdout(io.StringIO()),
+            ):
+                infer_main(
+                    [
+                        "--checkpoint",
+                        str(output / "checkpoint_last.pt"),
+                        "--output-dir",
+                        str(infer_output),
+                        "--sessions",
+                        "Session_4",
+                        "--device",
+                        "cpu",
+                        "--no-progress",
+                    ]
+                )
+            self.assertEqual(before_hash, sha256_file(output / "checkpoint_last.pt"))
+            np.testing.assert_array_equal(
+                np.load(output / "raw_eval" / "Session_4.npy"),
+                np.load(infer_output / "raw_eval" / "Session_4.npy"),
+            )
+            self.assertTrue(
+                json.loads((infer_output / "metrics.json").read_text())[
+                    "inference_only"
+                ]
+            )
            for method in (
                output,
                output / "stage_one_reference",