408 lines
16 KiB
Python
408 lines
16 KiB
Python
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()
|