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()