Add causal person track GRU residual experiment
This commit is contained in:
Xin Wang
+157
-2
@@ -33,7 +33,8 @@ from extract_turn_local_features import (
|
||||
)
|
||||
from local_turn_features import (
|
||||
LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes,
|
||||
load_frame_set_features, load_local_features, local_head_dimensions,
|
||||
load_frame_set_features, load_local_features, load_track_features,
|
||||
local_head_dimensions,
|
||||
pooled_recent_visual, reference_alignment,
|
||||
)
|
||||
from train_turn_head import _load_examples
|
||||
@@ -352,6 +353,72 @@ class LocalTests(unittest.TestCase):
|
||||
dataset="easycom",
|
||||
)
|
||||
|
||||
def test_track_loader_builds_causal_histories_and_resets_tracks(self):
|
||||
reference, side, _ = fixture(n=12)
|
||||
frame_times = torch.tensor([0., .5, 1., 1.5, 2.], dtype=torch.float64)
|
||||
side["persons"] = {
|
||||
"features": torch.arange(1, 9, dtype=torch.float32)[:, None].repeat(1, 6),
|
||||
"frame_times": frame_times,
|
||||
"frame_offsets": torch.tensor([0, 2, 4, 6, 7, 8]),
|
||||
"records": [
|
||||
{"id": "A", "source_clip_index": 0},
|
||||
{"id": "B", "source_clip_index": 0},
|
||||
{"id": "B", "source_clip_index": 0},
|
||||
{"id": "A", "source_clip_index": 0},
|
||||
{"id": "A", "source_clip_index": 0},
|
||||
{"id": -1, "source_clip_index": 0},
|
||||
# The same ID in another retained source clip is a new track.
|
||||
{"id": "A", "source_clip_index": 1},
|
||||
{"id": "A", "source_clip_index": 1},
|
||||
],
|
||||
}
|
||||
with tempfile.TemporaryDirectory() as directory:
|
||||
path = Path(directory) / "Session_4.pt"
|
||||
torch.save(side, path)
|
||||
pooled, people, contract = load_track_features(
|
||||
path,
|
||||
reference,
|
||||
sample_id="Session_4",
|
||||
dataset="easycom",
|
||||
track_max_gap_seconds=.75,
|
||||
)
|
||||
self.assertEqual(pooled.shape, (12, 13))
|
||||
self.assertEqual(contract["modalities"], "av")
|
||||
self.assertEqual(
|
||||
people.previous_observation.tolist(),
|
||||
[-1, -1, 1, 0, 3, -1, -1, 6],
|
||||
)
|
||||
# Cutoff 1.2s selects frame 1.0s. A has observations 1,4,5;
|
||||
# anonymous -1 remains a one-observation track.
|
||||
history, mask = people.padded_history(
|
||||
torch.tensor([5]),
|
||||
max_observations=4,
|
||||
)
|
||||
self.assertEqual(mask.tolist(), [[[True, True, True, False],
|
||||
[True, False, False, False]]])
|
||||
self.assertEqual(history[0, 0, :, 0].tolist(), [1., 4., 5., 0.])
|
||||
self.assertEqual(history[0, 1, :, 0].tolist(), [6., 0., 0., 0.])
|
||||
# Cutoff 1.8s selects frame 1.5s; source clip reset prevents A's
|
||||
# earlier observations from crossing the boundary.
|
||||
history, mask = people.padded_history(
|
||||
torch.tensor([8]),
|
||||
max_observations=4,
|
||||
)
|
||||
self.assertEqual(mask.tolist(), [[[True, False, False, False]]])
|
||||
self.assertEqual(history[0, 0, :, 0].tolist(), [7., 0., 0., 0.])
|
||||
|
||||
broken = copy.deepcopy(side)
|
||||
broken["persons"]["records"][1]["id"] = "A"
|
||||
torch.save(broken, path)
|
||||
with self.assertRaisesRegex(ValueError, "Duplicate person/track ID"):
|
||||
load_track_features(
|
||||
path,
|
||||
reference,
|
||||
sample_id="Session_4",
|
||||
dataset="easycom",
|
||||
track_max_gap_seconds=.75,
|
||||
)
|
||||
|
||||
def test_frame_person_concatenation_selection_and_empty_padding(self):
|
||||
features = torch.tensor([[1., 2.], [3., 4.]])
|
||||
a = FramePersonFeatures(
|
||||
@@ -567,6 +634,86 @@ class LocalTests(unittest.TestCase):
|
||||
restored.eval()
|
||||
self.assertTrue(torch.equal(restored(states, local, people, mask), logits))
|
||||
|
||||
def test_track_gru_residual_starts_from_exact_pooled_logits(self):
|
||||
torch.manual_seed(17)
|
||||
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,
|
||||
}
|
||||
pooled = TurnPredictionHead(**common, local_fusion="pooled").eval()
|
||||
track = TurnPredictionHead(
|
||||
**common,
|
||||
local_fusion="track-gru-residual",
|
||||
person_attention_dim=4,
|
||||
person_attention_heads=2,
|
||||
person_history_frames=3,
|
||||
).eval()
|
||||
track_state = track.state_dict()
|
||||
for name, value in pooled.state_dict().items():
|
||||
self.assertIn(name, track_state)
|
||||
self.assertEqual(track_state[name].shape, value.shape)
|
||||
track_state[name] = value
|
||||
track.load_state_dict(track_state)
|
||||
|
||||
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)
|
||||
self.assertTrue(torch.equal(track(states, local, histories, mask), baseline))
|
||||
self.assertEqual(track.person_interaction_gates(), {"person_track_residual": 0.0})
|
||||
|
||||
with torch.no_grad():
|
||||
assert track.person_track_residual_gate is not None
|
||||
track.person_track_residual_gate.fill_(.25)
|
||||
logits = track(states, local, histories, mask)
|
||||
self.assertFalse(torch.equal(logits[:2], baseline[:2]))
|
||||
# With no current person the correction must be exactly zero.
|
||||
self.assertTrue(torch.equal(logits[2], baseline[2]))
|
||||
permutation = torch.tensor([1, 0])
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
logits,
|
||||
track(states, local, histories[:, permutation], mask[:, permutation]),
|
||||
atol=1e-6,
|
||||
rtol=1e-6,
|
||||
)
|
||||
)
|
||||
changed_padding = histories.clone()
|
||||
changed_padding[~mask] = 1000
|
||||
self.assertTrue(torch.equal(logits, track(states, local, changed_padding, mask)))
|
||||
logits.square().mean().backward()
|
||||
self.assertIsNotNone(track.person_track_gru.weight_ih_l0.grad)
|
||||
self.assertIsNotNone(track.person_track_delta[-1].weight.grad)
|
||||
self.assertIsNotNone(track.person_track_residual_gate.grad)
|
||||
|
||||
config = {
|
||||
**common,
|
||||
"local_fusion": "track-gru-residual",
|
||||
"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(track.state_dict())
|
||||
self.assertTrue(torch.equal(restored(states, local, histories, 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)):
|
||||
@@ -855,6 +1002,7 @@ class RoundTripTests(unittest.TestCase):
|
||||
"audio",
|
||||
"visual",
|
||||
"av",
|
||||
"track-gru",
|
||||
"frame-set",
|
||||
"f5-screen",
|
||||
"f6-screen",
|
||||
@@ -877,6 +1025,7 @@ class RoundTripTests(unittest.TestCase):
|
||||
if mode == "av":
|
||||
train += ["--transition-objective", "r1-supervised"]
|
||||
person_mode = {
|
||||
"track-gru": "track-gru-residual",
|
||||
"frame-set": "frame-set-attention",
|
||||
"f5-screen": "audio-person-frame-set",
|
||||
"f6-screen": "gated-person-state-adapter",
|
||||
@@ -887,10 +1036,16 @@ class RoundTripTests(unittest.TestCase):
|
||||
"--person-attention-dim", "4",
|
||||
"--person-attention-heads", "2",
|
||||
]
|
||||
if mode == "track-gru":
|
||||
train += [
|
||||
"--person-history-frames", "3",
|
||||
"--person-track-max-gap-seconds", "1",
|
||||
"--initialize-pooled-checkpoint", str(root / "av.pth"),
|
||||
]
|
||||
if mode != "baseline":
|
||||
modality = (
|
||||
"av"
|
||||
if mode in ("ego4d", "frame-set", "f5-screen", "f6-screen")
|
||||
if mode in ("ego4d", "track-gru", "frame-set", "f5-screen", "f6-screen")
|
||||
else mode
|
||||
)
|
||||
train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality]
|
||||
|
||||
Reference in New Issue
Block a user