wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add selective EgoSpeak Omni stage-two finetuning

Browse Source
This commit is contained in:
Xin Wang
2026-09-16 17:04:46 +08:00
parent eff7a5ad29
commit 64adbef13d
7 changed files with 1298 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

374
test_egospeak_omni_finetune.py Normal file
View File

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