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

feat: add causal frame set attention fusion

Browse Source
This commit is contained in:
Xin Wang
2026-09-02 18:48:26 +08:00
parent e769094b3d
commit 8e70281330
8 changed files with 837 additions and 38 deletions
Download Patch File Download Diff File Expand all files Collapse all files

174
test_local_turn_features.py
View File

@@ -28,8 +28,9 @@ from extract_turn_local_features import (
main as extract_main,
)
from local_turn_features import (
LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes, load_local_features,
local_head_dimensions, pooled_recent_visual, reference_alignment,
LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes,
load_frame_set_features, load_local_features, local_head_dimensions,
pooled_recent_visual, reference_alignment,
)
from train_turn_head import _load_examples
@@ -47,6 +48,17 @@ def fixture(n=20, name="Session_4", dataset="easycom"):
"causal_policy": "half-open-trailing-window", "uses_speech_labels": False}}
side["visual_valid"][0] = False
side["visual"][0] = 0
frame_times = torch.arange(max(1, (n + 2) // 3), dtype=torch.float64) * .5
counts = [2 if index % 3 == 0 else (0 if index % 3 == 1 else 1)
for index in range(len(frame_times))]
offsets = torch.tensor([0, *np.cumsum(counts).tolist()], dtype=torch.long)
person_features = torch.randn(int(offsets[-1]), 6)
side["persons"] = {
"features": person_features,
"frame_times": frame_times,
"frame_offsets": offsets,
"records": [{"id": str(index)} for index in range(len(person_features))],
}
labels = np.tile(np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 1, 1]], dtype=np.float32), (n // 4 + 1, 1))[:n]
return reference, side, labels
@@ -280,6 +292,148 @@ class LocalTests(unittest.TestCase):
with self.assertRaisesRegex(ValueError, "times"):
reference_alignment({**reference, "times": reference["times"] + .2})
def test_frame_set_loader_is_causal_ragged_and_fail_closed(self):
reference, side, _ = fixture()
with tempfile.TemporaryDirectory() as directory:
path = Path(directory) / "Session_4.pt"
torch.save(side, path)
audio, people, contract = load_frame_set_features(
path,
reference,
sample_id="Session_4",
dataset="easycom",
)
self.assertEqual(audio.shape, (20, 6))
self.assertEqual(contract["modalities"], "av")
self.assertEqual(people.dimension, 6)
self.assertTrue(
torch.all(
people.chosen_frame_times[
torch.isfinite(people.chosen_frame_times)
]
< reference["times"][
torch.isfinite(people.chosen_frame_times)
]
)
)
# At 0.2/0.4s, both rows reuse the 0.0s frame with two people.
values, mask = people.padded(torch.tensor([0, 1, 2]))
self.assertEqual(values.shape, (3, 2, 6))
self.assertEqual(mask.tolist(), [[True, True], [True, True], [False, False]])
self.assertTrue(torch.equal(values[0], values[1]))
self.assertEqual(torch.count_nonzero(values[2]), 0)
# The frame stamped exactly 1.0s is excluded at the 1.0s cutoff.
self.assertEqual(people.chosen_frame_times[4].item(), .5)
broken_cases = []
broken = copy.deepcopy(side)
broken["persons"]["frame_offsets"][-1] -= 1
broken_cases.append(broken)
broken = copy.deepcopy(side)
broken["persons"]["frame_times"][1] = broken["persons"]["frame_times"][0]
broken_cases.append(broken)
broken = copy.deepcopy(side)
broken["persons"]["records"] = []
broken_cases.append(broken)
broken = copy.deepcopy(side)
broken["persons"]["features"][0, 0] = float("nan")
broken_cases.append(broken)
for index, broken in enumerate(broken_cases):
torch.save(broken, path)
with self.subTest(index=index), self.assertRaises(ValueError):
load_frame_set_features(
path,
reference,
sample_id="Session_4",
dataset="easycom",
)
def test_frame_person_concatenation_selection_and_empty_padding(self):
features = torch.tensor([[1., 2.], [3., 4.]])
a = FramePersonFeatures(
features,
torch.tensor([0, 2], dtype=torch.long),
torch.tensor([0, -1], dtype=torch.long),
torch.tensor([0., float("nan")], dtype=torch.float64),
)
b = FramePersonFeatures(
torch.empty(0, 2),
torch.tensor([0, 0], dtype=torch.long),
torch.tensor([0], dtype=torch.long),
torch.tensor([.5], dtype=torch.float64),
)
combined = FramePersonFeatures.concatenate([a, b])
self.assertEqual(len(combined), 3)
self.assertEqual(combined.row_to_frame.tolist(), [0, -1, 1])
selected = combined.select_rows(torch.tensor([2, 0]))
values, mask = selected.padded()
self.assertEqual(mask.tolist(), [[False, False], [True, True]])
self.assertEqual(torch.count_nonzero(values[0]), 0)
self.assertTrue(torch.equal(values[1], features))
def test_frame_set_head_is_permutation_invariant_and_masks_padding(self):
torch.manual_seed(7)
head = TurnPredictionHead(
hidden_size=8,
num_input_layers=4,
layer_fusion="concat-mlp",
local_visual_dim=6,
local_audio_dim=5,
local_projection_dim=3,
local_fusion="frame-set-attention",
person_attention_dim=4,
person_attention_heads=2,
transition_auxiliary=True,
dropout=0,
).eval()
states = torch.randn(3, 4, 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, transitions = head.forward_with_transition(
states,
local,
people,
mask,
)
self.assertEqual(logits.shape, (3, 10, 3))
self.assertEqual(transitions.shape, (3, 10))
(logits.square().mean() + transitions.square().mean()).backward()
self.assertIsNotNone(head.person_attention.in_proj_weight.grad)
permutation = torch.tensor([1, 0, 2])
permuted = head(states, local, people[:, permutation], mask[:, permutation])
self.assertTrue(torch.allclose(logits, permuted, atol=1e-6, rtol=1e-6))
changed = people.clone()
changed[~mask] = 1000
self.assertTrue(torch.allclose(logits, head(states, local, changed, mask)))
with self.assertRaisesRegex(ValueError, "person features"):
head(states, local)
config = {
"hidden_size": 8,
"num_input_layers": 4,
"layer_fusion": "concat-mlp",
"local_visual_dim": 6,
"local_audio_dim": 5,
"local_projection_dim": 3,
"local_fusion": "frame-set-attention",
"person_attention_dim": 4,
"person_attention_heads": 2,
"transition_auxiliary": True,
"dropout": 0,
}
restored = TurnPredictionHead(
**checkpoint_head_config({"head_config": config})
)
restored.load_state_dict(head.state_dict())
restored.eval()
self.assertTrue(torch.equal(restored(states, local, people, mask), logits))
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)):
@@ -338,7 +492,8 @@ class LocalTests(unittest.TestCase):
np.save(root / "targets/Session_4.npy", targets)
data = _load_examples(root / "features", root / "targets", ["Session_4"], "open-stream-replay",
"easycom", {"Session_4": {1, 4}}, root / "local", "audio")
self.assertTrue(torch.equal(data[-1][:, :5], side["audio"][[1, 4]]))
self.assertTrue(torch.equal(data[-2][:, :5], side["audio"][[1, 4]]))
self.assertIsNone(data[-1])
self.assertTrue(torch.equal(data[0], reference["states"][[1, 4]]))
self.assertTrue(torch.equal(data[1][:, 0], torch.from_numpy(targets[[2, 5]])))
@@ -562,7 +717,7 @@ class RoundTripTests(unittest.TestCase):
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"]
for mode in ("baseline", "audio", "visual", "av", "ego4d"):
for mode in ("baseline", "audio", "visual", "av", "frame-set", "ego4d"):
checkpoint = root / f"{mode}.pth"
train = [sys.executable, str(repo / "train_turn_head.py"), *common,
"--epochs", "1", "--warmup-epochs", "0", "--batch-size", "4",
@@ -579,8 +734,15 @@ class RoundTripTests(unittest.TestCase):
train += ["--dev-sessions", "5", "--local-projection-dim", "0"]
if mode == "av":
train += ["--transition-objective", "r1-supervised"]
if mode == "frame-set":
train += [
"--local-fusion", "frame-set-attention",
"--person-attention-dim", "4",
"--person-attention-heads", "2",
]
if mode != "baseline":
train += ["--local-feature-root", str(root / "local"), "--local-modalities", "av" if mode == "ego4d" else mode]
modality = "av" if mode in ("ego4d", "frame-set") else mode
train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality]
infer += ["--local-feature-root", str(root / "local")]
for command in (train, infer):
completed = subprocess.run(command, cwd=repo, env=env, capture_output=True, text=True, timeout=60)
@@ -590,6 +752,8 @@ class RoundTripTests(unittest.TestCase):
self.assertEqual(bool(metrics["local_contract"]), mode != "baseline")
if mode == "av":
self.assertTrue(metrics["transition_auxiliary"])
if mode == "frame-set":
self.assertEqual(metrics["local_fusion"], "frame-set-attention")
# The reporting command uses the same metric keys as inference.
completed = subprocess.run(
[sys.executable, str(repo / "summarize_local_baselines.py"), "--root", str(root)],