feat: add causal frame set attention fusion
This commit is contained in:
Xin Wang
+169
-5
@@ -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)],
|
||||
|
||||
Reference in New Issue
Block a user