Add causal person fusion and MuVAP stage-3 baselines

This commit is contained in:
Xin Wang
2026-09-10 06:54:22 +08:00
parent 0534c9d3d0
commit d426b4740a
225 changed files with 27716 additions and 119 deletions

View File

@@ -31,9 +31,18 @@ from extract_turn_local_features import (
TESTED_TRANSFORMERS_VERSIONS, validate_local_transformers_version,
main as extract_main,
)
from extract_highrate_person_av_features import (
MODEL_FPS as HIGHRATE_MODEL_FPS,
_chunked_audio_encode,
_chunked_person_encode,
_model_tick_bounds,
_square_gray_face,
)
from local_turn_features import (
LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes,
load_frame_set_features, load_local_features, load_track_features,
HIGHRATE_PERSON_AV_FORMAT, LOCAL_FORMAT, FramePersonFeatures,
causal_frame_ticks, crop_box, easycom_boxes,
load_frame_set_features, load_highrate_person_av_features,
load_local_features, load_track_features,
local_head_dimensions,
pooled_recent_visual, reference_alignment,
)
@@ -68,6 +77,42 @@ def fixture(n=20, name="Session_4", dataset="easycom"):
return reference, side, labels
def highrate_fixture(n=20, name="Session_4", dataset="easycom"):
reference, side, labels = fixture(n=n, name=name, dataset=dataset)
side["format"] = HIGHRATE_PERSON_AV_FORMAT
side["visual"] = torch.randn(n, 6)
side["visual_valid"] = torch.ones(n, dtype=torch.bool)
side["persons"] = {
"features": side["visual"].clone(),
"frame_times": reference["times"].clone(),
"frame_offsets": torch.arange(n + 1, dtype=torch.long),
"records": [
{"id": "p", "asd_interval_end": float(reference["times"][row])}
for row in range(n)
],
}
side["config"].update(
{
"model_fps": 25.0,
"output_hz": 5.0,
"person_feature_kind": (
"muvap-asd-causal-visual-plus-activity-logit"
),
"person_visual_dim": 5,
"person_activity_logit_dim": 1,
"audio_feature_kind": "muvap-cpc-causal-audio-state",
"audio_source": "single-shared-traditional-cpc-stream",
"person_tokens_duplicate_audio": False,
"asd_audio_preprocessing": (
"causal-float-pcm-clip-only-no-rms-normalization"
),
"asd_temporal_attention": "causal",
"causal_policy": "right-edge-causal-highrate-asd",
}
)
return reference, side, labels
class LocalTests(unittest.TestCase):
def setUp(self):
torch.set_num_threads(1)
@@ -94,6 +139,110 @@ class LocalTests(unittest.TestCase):
self.assertFalse(present)
self.assertEqual(value.item(), 0.)
def test_highrate_person_av_right_edges_are_causal_and_inclusive(self):
reference, side, _ = highrate_fixture(n=5)
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "Session_4.pt"
torch.save(side, path)
local, people, contract = load_highrate_person_av_features(
path,
reference,
sample_id="Session_4",
dataset="easycom",
)
packed_audio, audio_contract = load_local_features(
path,
reference,
sample_id="Session_4",
dataset="easycom",
modalities="audio",
)
self.assertEqual(local.shape, (5, 6))
self.assertEqual(people.row_to_frame.tolist(), [0, 1, 2, 3, 4])
self.assertTrue(torch.equal(people.chosen_frame_times, reference["times"]))
self.assertEqual(contract["format"], HIGHRATE_PERSON_AV_FORMAT)
self.assertEqual(packed_audio.shape, (5, 6))
self.assertEqual(audio_contract["modalities"], "audio")
values, mask = people.padded()
self.assertEqual(values.shape, (5, 1, 6))
self.assertTrue(mask.all())
def test_highrate_person_av_rejects_old_or_duplicated_audio_contracts(self):
reference, side, _ = highrate_fixture(n=5)
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "Session_4.pt"
old = copy.deepcopy(side)
old["format"] = "causal-highrate-person-av-v1"
torch.save(old, path)
with self.assertRaisesRegex(ValueError, "Invalid local feature format"):
load_highrate_person_av_features(
path,
reference,
sample_id="Session_4",
dataset="easycom",
)
duplicated = copy.deepcopy(side)
duplicated["config"]["person_tokens_duplicate_audio"] = True
torch.save(duplicated, path)
with self.assertRaisesRegex(ValueError, "causal ASD contract"):
load_highrate_person_av_features(
path,
reference,
sample_id="Session_4",
dataset="easycom",
)
def test_highrate_face_geometry_and_chunked_causal_coverage(self):
image = Image.new("RGB", (40, 20), "white")
face = _square_gray_face(image)
self.assertEqual(face.shape, (112, 112))
self.assertEqual(face.dtype, np.uint8)
self.assertEqual(int(face[0, 0]), 128)
clip = SimpleNamespace(dense_start=0.0, dense_end=1.0)
self.assertEqual(_model_tick_bounds(clip, 1.0), (0, 25))
self.assertEqual(HIGHRATE_MODEL_FPS, 25.0)
class FakeEncoder:
audio_dim = 2
feature_dim = 3
def __init__(self):
self.audio_lengths = []
self.person_lengths = []
def encode_audio(self, waveform, num_frames):
self.audio_lengths.append(num_frames)
return torch.stack(
(torch.arange(num_frames), torch.ones(num_frames)), dim=-1
).float()
def encode_person(self, faces, audio_states):
self.person_lengths.append(len(faces))
return torch.cat((audio_states, torch.ones(len(faces), 1)), dim=-1)
encoder = FakeEncoder()
faces = np.zeros((9, 112, 112), dtype=np.uint8)
waveform = np.zeros(9 * 640, dtype=np.float32)
audio_result = _chunked_audio_encode(
encoder,
waveform,
num_frames=len(faces),
output_frames=4,
context_frames=2,
)
person_result = _chunked_person_encode(
encoder,
faces,
audio_result,
output_frames=4,
context_frames=2,
)
self.assertEqual(audio_result.shape, (9, 2))
self.assertEqual(person_result.shape, (9, 3))
self.assertEqual(encoder.audio_lengths, [4, 6, 3])
self.assertEqual(encoder.person_lengths, [4, 6, 3])
def test_box_formats_and_crop(self):
person = {"Participant_ID": -1, "x1": 2, "y1": 3, "x2": 7, "y2": 9}
faces = easycom_boxes([{"Frame_Number": 1, "Participants": [person]}], "face")
@@ -505,6 +654,95 @@ class LocalTests(unittest.TestCase):
restored.eval()
self.assertTrue(torch.equal(restored(states, local, people, mask), logits))
def test_highrate_person_av_starts_from_exact_audio_only_base(self):
torch.manual_seed(19)
audio = TurnPredictionHead(
hidden_size=8,
local_visual_dim=0,
local_audio_dim=5,
local_projection_dim=3,
local_fusion="pooled",
person_attention_dim=4,
person_attention_heads=2,
dropout=0,
).eval()
highrate = TurnPredictionHead(
hidden_size=8,
local_visual_dim=6,
local_audio_dim=5,
local_projection_dim=3,
local_fusion="highrate-person-av-attention",
person_attention_dim=4,
person_attention_heads=2,
dropout=0,
).eval()
highrate_state = highrate.state_dict()
for name, value in audio.state_dict().items():
self.assertIn(name, highrate_state)
self.assertEqual(tuple(value.shape), tuple(highrate_state[name].shape))
highrate_state[name] = value
highrate.load_state_dict(highrate_state)
states = torch.randn(4, 8)
local = torch.randn(4, 6)
local[:, -1] = 1
people = torch.randn(4, 3, 6)
mask = torch.tensor(
[[True, True, False], [True, False, False], [False, False, False],
[True, True, True]]
)
people[~mask] = 0
reference = audio(states, local)
actual = highrate(states, local, people, mask)
self.assertTrue(torch.equal(reference, actual))
diagnostics = highrate.person_adapter_diagnostics()
self.assertEqual(diagnostics["residual_output_weight_norm"], 0.0)
self.assertEqual(diagnostics["delta_logit_max_abs"], 0.0)
actual.sum().backward()
assert highrate.highrate_person_residual is not None
self.assertGreater(
float(highrate.highrate_person_residual.weight.grad.abs().sum()),
0.0,
)
def test_highrate_person_av_head_uses_variable_person_sets(self):
torch.manual_seed(9)
head = TurnPredictionHead(
hidden_size=8,
local_visual_dim=6,
local_audio_dim=5,
local_projection_dim=3,
local_fusion="highrate-person-av-attention",
person_attention_dim=4,
person_attention_heads=2,
dropout=0,
).eval()
states = torch.randn(3, 8)
local = torch.randn(3, 6)
local[:, -1] = 1
people = torch.randn(3, 3, 6)
mask = torch.tensor(
[[True, True, False], [True, False, False], [False, False, False]]
)
people[~mask] = 0
logits = head(states, local, people, mask)
self.assertEqual(logits.shape, (3, 10, 3))
permutation = torch.tensor([1, 0, 2])
self.assertTrue(
torch.allclose(
logits,
head(states, local, people[:, permutation], mask[:, permutation]),
atol=1e-6,
rtol=1e-6,
)
)
changed = people.clone()
changed[~mask] = 1000
self.assertTrue(torch.equal(logits, head(states, local, changed, mask)))
logits.square().mean().backward()
self.assertIsNotNone(head.person_attention.in_proj_weight.grad)
def test_f5_screening_audio_person_read_is_gated_and_permutation_invariant(self):
torch.manual_seed(11)
common = {
@@ -714,6 +952,227 @@ class LocalTests(unittest.TestCase):
restored.load_state_dict(track.state_dict())
self.assertTrue(torch.equal(restored(states, local, histories, mask), logits))
def test_causal_person_adapter_v2_learns_after_exact_pooled_start(self):
common = {
"hidden_size": 8,
"num_input_layers": 4,
"layer_fusion": "concat-mlp",
"local_visual_dim": 6,
"local_audio_dim": 5,
"local_projection_dim": 3,
"dropout": 0,
}
for mode in (
"causal-person-adapter-v2",
"causal-person-adapter-egospeak-gru",
):
with self.subTest(mode=mode):
torch.manual_seed(23)
pooled = TurnPredictionHead(
**common,
local_fusion="pooled",
).eval()
adapter = TurnPredictionHead(
**common,
local_fusion=mode,
person_attention_dim=4,
person_attention_heads=2,
person_history_frames=3,
).eval()
adapter_state = adapter.state_dict()
for name, value in pooled.state_dict().items():
self.assertIn(name, adapter_state)
self.assertEqual(adapter_state[name].shape, value.shape)
adapter_state[name] = value
adapter.load_state_dict(adapter_state)
assert adapter.person_track_residual_output is not None
self.assertEqual(
torch.count_nonzero(
adapter.person_track_residual_output.weight
),
0,
)
self.assertEqual(
torch.count_nonzero(
adapter.person_track_residual_output.bias
),
0,
)
frozen_parameters = {}
for name, parameter in adapter.named_parameters():
if name == "layer_logits" or name.startswith(
("layer_fusion_net.", "local_projections.", "net.")
):
parameter.requires_grad_(False)
frozen_parameters[name] = parameter.detach().clone()
states = torch.randn(3, 4, 8)
local = torch.randn(3, 13)
local[:, 6] = 1
local[:, 12] = 1
histories = torch.randn(3, 2, 3, 6)
mask = torch.tensor(
[
[[True, True, True], [True, True, False]],
[[True, False, False], [False, False, False]],
[[False, False, False], [False, False, False]],
]
)
histories[~mask] = 0
baseline = pooled(states, local)
initial = adapter(states, local, histories, mask)
self.assertTrue(torch.equal(initial, baseline))
self.assertEqual(
adapter.person_adapter_diagnostics()["delta_logit_max_abs"],
0.0,
)
optimizer = torch.optim.SGD(adapter.parameters(), lr=0.1)
target = torch.randn_like(initial)
(initial - target).square().mean().backward()
output_gradient = (
adapter.person_track_residual_output.weight.grad
)
self.assertIsNotNone(output_gradient)
self.assertGreater(int(torch.count_nonzero(output_gradient)), 0)
optimizer.step()
optimizer.zero_grad(set_to_none=True)
trained = adapter(states, local, histories, mask)
(trained - target).square().mean().backward()
input_projection = (
adapter.person_track_input_embedding[0].weight
if adapter.person_track_input_embedding is not None
else adapter.person_input_projection[1].weight
)
for parameter in (
input_projection,
adapter.person_track_gru.weight_ih_l0,
adapter.person_track_gru.weight_hh_l0,
adapter.person_attention.in_proj_weight,
adapter.person_track_local_state[1].weight,
adapter.person_track_dynamic_gate[1].weight,
):
self.assertIsNotNone(parameter.grad)
self.assertTrue(torch.isfinite(parameter.grad).all())
self.assertGreater(int(torch.count_nonzero(parameter.grad)), 0)
for name, expected in frozen_parameters.items():
self.assertTrue(
torch.equal(dict(adapter.named_parameters())[name], expected),
name,
)
# Empty rows remain exactly on the pooled path after the
# residual opens.
self.assertTrue(
torch.equal(trained[2], pooled(states, local)[2])
)
permutation = torch.tensor([1, 0])
self.assertTrue(
torch.allclose(
trained,
adapter(
states,
local,
histories[:, permutation],
mask[:, permutation],
),
atol=1e-6,
rtol=1e-6,
)
)
changed_padding = histories.clone()
changed_padding[~mask] = 1000
self.assertTrue(
torch.equal(
trained,
adapter(states, local, changed_padding, mask),
)
)
config = {
**common,
"local_fusion": mode,
"person_attention_dim": 4,
"person_attention_heads": 2,
"person_history_frames": 3,
}
restored = TurnPredictionHead(
**checkpoint_head_config({"head_config": config})
).eval()
restored.load_state_dict(adapter.state_dict())
self.assertTrue(
torch.equal(
restored(states, local, histories, mask),
adapter(states, local, histories, mask),
)
)
def test_egospeak_style_track_gru_matches_explicit_reference(self):
torch.manual_seed(29)
head = TurnPredictionHead(
hidden_size=8,
local_visual_dim=6,
local_audio_dim=5,
local_projection_dim=3,
local_fusion="causal-person-adapter-egospeak-gru",
person_attention_dim=4,
person_attention_heads=2,
person_history_frames=3,
dropout=0,
).eval()
embedding = head.person_track_input_embedding
self.assertIsNotNone(embedding)
self.assertEqual(
[type(module) for module in embedding],
[
torch.nn.Linear,
torch.nn.LayerNorm,
torch.nn.ReLU,
torch.nn.Dropout,
],
)
self.assertEqual(head.person_track_gru.num_layers, 1)
self.assertFalse(head.person_track_gru.bidirectional)
self.assertTrue(head.person_track_gru.batch_first)
self.assertIsInstance(
head.person_track_output_activation,
torch.nn.ReLU,
)
histories = torch.randn(2, 2, 3, 6)
mask = torch.tensor(
[
[[True, True, True], [True, True, False]],
[[True, False, False], [False, False, False]],
]
)
histories[~mask] = 1234
actual, valid = head._encode_person_track_histories(histories, mask)
expected = torch.zeros_like(actual)
for batch in range(2):
for person in range(2):
length = int(mask[batch, person].sum())
if not length:
continue
sequence = embedding(
histories[batch, person, :length]
)[None]
initial = sequence.new_zeros(1, 1, 4)
_, hidden = head.person_track_gru(sequence, initial)
expected[batch, person] = torch.relu(hidden[-1, 0])
self.assertTrue(torch.equal(valid, mask.any(dim=-1)))
self.assertTrue(torch.allclose(actual, expected, atol=1e-7, rtol=0))
self.assertGreaterEqual(float(actual.min()), 0.0)
changed_padding = histories.clone()
changed_padding[~mask] = -9876
changed, _ = head._encode_person_track_histories(
changed_padding,
mask,
)
self.assertTrue(torch.equal(actual, changed))
def test_head_shapes_gradients_missing_and_roundtrip(self):
states = torch.randn(3, 4, 8)
for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)):
@@ -983,7 +1442,7 @@ class RoundTripTests(unittest.TestCase):
env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1")
with tempfile.TemporaryDirectory() as directory:
root = Path(directory)
for folder in ("features", "local", "targets"):
for folder in ("features", "local", "highrate", "targets"):
(root / folder).mkdir()
train_uid = "00000000-0000-0000-0000-000000000004"
test_uid = "00000000-0000-0000-0000-000000000001"
@@ -994,6 +1453,11 @@ class RoundTripTests(unittest.TestCase):
reference, side, targets = fixture(name=name, dataset=dataset)
torch.save(reference, root / "features" / f"{name}.pt")
torch.save(side, root / "local" / f"{name}.pt")
_, highrate_side, _ = highrate_fixture(
name=name,
dataset=dataset,
)
torch.save(highrate_side, root / "highrate" / f"{name}.pt")
np.save(root / "targets" / f"{name}.npy", targets)
common = ["--feature-root", str(root / "features"), "--target-root", str(root / "targets"),
"--device", "cpu", "--no-progress", "--expected-mode", "open-stream-replay"]
@@ -1003,9 +1467,13 @@ class RoundTripTests(unittest.TestCase):
"visual",
"av",
"track-gru",
"person-adapter-v2",
"egospeak-person-gru",
"frame-set",
"f5-screen",
"f6-screen",
"traditional-audio",
"highrate-person-av",
"ego4d",
):
checkpoint = root / f"{mode}.pth"
@@ -1026,9 +1494,14 @@ class RoundTripTests(unittest.TestCase):
train += ["--transition-objective", "r1-supervised"]
person_mode = {
"track-gru": "track-gru-residual",
"person-adapter-v2": "causal-person-adapter-v2",
"egospeak-person-gru": (
"causal-person-adapter-egospeak-gru"
),
"frame-set": "frame-set-attention",
"f5-screen": "audio-person-frame-set",
"f6-screen": "gated-person-state-adapter",
"highrate-person-av": "highrate-person-av-attention",
}.get(mode)
if person_mode is not None:
train += [
@@ -1036,20 +1509,42 @@ class RoundTripTests(unittest.TestCase):
"--person-attention-dim", "4",
"--person-attention-heads", "2",
]
if mode == "track-gru":
if mode in (
"track-gru",
"person-adapter-v2",
"egospeak-person-gru",
):
train += [
"--person-history-frames", "3",
"--person-track-max-gap-seconds", "1",
"--initialize-pooled-checkpoint", str(root / "av.pth"),
"--transition-objective", "r1-supervised",
]
if mode == "highrate-person-av":
train += [
"--initialize-pooled-checkpoint",
str(root / "traditional-audio.pth"),
]
if mode != "baseline":
modality = (
"av"
if mode in ("ego4d", "track-gru", "frame-set", "f5-screen", "f6-screen")
if mode in (
"ego4d", "track-gru", "person-adapter-v2",
"egospeak-person-gru",
"frame-set", "f5-screen", "f6-screen",
"highrate-person-av",
)
else "audio"
if mode == "traditional-audio"
else mode
)
train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality]
infer += ["--local-feature-root", str(root / "local")]
local_root = root / (
"highrate"
if mode in ("traditional-audio", "highrate-person-av")
else "local"
)
train += ["--local-feature-root", str(local_root), "--local-modalities", modality]
infer += ["--local-feature-root", str(local_root)]
for command in (train, infer):
completed = subprocess.run(command, cwd=repo, env=env, capture_output=True, text=True, timeout=60)
self.assertEqual(completed.returncode, 0, f"{mode}: {completed.stdout}\n{completed.stderr}")