icassp2027-experiment/test_egospeak_omni_finetune.py

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,
    sha256_file,
)
from infer_egospeak_omni_finetune import main as infer_main
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())
            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",
                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()