feat: add person fusion screening variants
This commit is contained in:
Xin Wang
@@ -18,7 +18,11 @@ import numpy as np
|
||||
import torch
|
||||
from PIL import Image, ImageDraw
|
||||
|
||||
from easycom_head import TurnPredictionHead, checkpoint_head_config
|
||||
from easycom_head import (
|
||||
LOCAL_FUSION_SCOPES,
|
||||
TurnPredictionHead,
|
||||
checkpoint_head_config,
|
||||
)
|
||||
from extract_turn_local_features import (
|
||||
FrozenQwenLocalEncoder, ego4d_clip_boxes, extract_local,
|
||||
validate_source_timeline,
|
||||
@@ -434,6 +438,135 @@ class LocalTests(unittest.TestCase):
|
||||
restored.eval()
|
||||
self.assertTrue(torch.equal(restored(states, local, people, mask), logits))
|
||||
|
||||
def test_f5_screening_audio_person_read_is_gated_and_permutation_invariant(self):
|
||||
torch.manual_seed(11)
|
||||
common = {
|
||||
"hidden_size": 8,
|
||||
"local_visual_dim": 6,
|
||||
"local_audio_dim": 5,
|
||||
"local_projection_dim": 3,
|
||||
"person_attention_dim": 4,
|
||||
"person_attention_heads": 2,
|
||||
"dropout": 0,
|
||||
}
|
||||
f1 = TurnPredictionHead(
|
||||
**common,
|
||||
local_fusion="frame-set-attention",
|
||||
).eval()
|
||||
f5 = TurnPredictionHead(
|
||||
**common,
|
||||
local_fusion="audio-person-frame-set",
|
||||
).eval()
|
||||
# Align all shared parameters. A zero F5 gate must reproduce F1
|
||||
# exactly, independent of the newly initialized reverse-read block.
|
||||
f5_state = f5.state_dict()
|
||||
for name, value in f1.state_dict().items():
|
||||
if name in f5_state and f5_state[name].shape == value.shape:
|
||||
f5_state[name] = value
|
||||
f5.load_state_dict(f5_state)
|
||||
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
|
||||
baseline = f1(states, local, people, mask)
|
||||
self.assertTrue(torch.equal(f5(states, local, people, mask), baseline))
|
||||
self.assertEqual(f5.person_interaction_gates(), {"audio_person": 0.0})
|
||||
|
||||
with torch.no_grad():
|
||||
assert f5.audio_person_gate is not None
|
||||
f5.audio_person_gate.fill_(0.25)
|
||||
logits = f5(states, local, people, mask)
|
||||
self.assertFalse(torch.equal(logits, baseline))
|
||||
permutation = torch.tensor([1, 0, 2])
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
logits,
|
||||
f5(states, local, people[:, permutation], mask[:, permutation]),
|
||||
atol=1e-6,
|
||||
rtol=1e-6,
|
||||
)
|
||||
)
|
||||
logits.square().mean().backward()
|
||||
self.assertIsNotNone(f5.audio_person_attention.in_proj_weight.grad)
|
||||
self.assertIsNotNone(f5.audio_person_gate.grad)
|
||||
|
||||
def test_f6_screening_adapter_is_zero_gated_and_visual_only_when_valid(self):
|
||||
torch.manual_seed(13)
|
||||
head = TurnPredictionHead(
|
||||
hidden_size=8,
|
||||
local_visual_dim=6,
|
||||
local_audio_dim=5,
|
||||
local_projection_dim=3,
|
||||
local_fusion="gated-person-state-adapter",
|
||||
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, 2, 6)
|
||||
mask = torch.tensor([[True, True], [True, False], [False, False]])
|
||||
people[~mask] = 0
|
||||
zero_gate = head(states, local, people, mask)
|
||||
changed = people.clone()
|
||||
changed[mask] += torch.arange(6, dtype=changed.dtype) * 10
|
||||
# Before training, person memory is an exact zero residual.
|
||||
self.assertTrue(torch.equal(head(states, local, changed, mask), zero_gate))
|
||||
self.assertEqual(
|
||||
head.person_interaction_gates(),
|
||||
{"person_state_adapter": 0.0},
|
||||
)
|
||||
|
||||
with torch.no_grad():
|
||||
assert head.person_state_adapter_gate is not None
|
||||
head.person_state_adapter_gate.fill_(0.25)
|
||||
logits = head(states, local, people, mask)
|
||||
self.assertFalse(torch.equal(head(states, local, changed, mask), logits))
|
||||
# An empty set must remain visually inert even after opening the gate.
|
||||
empty_changed = people.clone()
|
||||
empty_changed[2] = 1000
|
||||
self.assertTrue(
|
||||
torch.equal(
|
||||
logits[2],
|
||||
head(states, local, empty_changed, mask)[2],
|
||||
)
|
||||
)
|
||||
permutation = torch.tensor([1, 0])
|
||||
self.assertTrue(
|
||||
torch.allclose(
|
||||
logits,
|
||||
head(states, local, people[:, permutation], mask[:, permutation]),
|
||||
atol=1e-6,
|
||||
rtol=1e-6,
|
||||
)
|
||||
)
|
||||
logits.square().mean().backward()
|
||||
self.assertIsNotNone(head.person_attention.in_proj_weight.grad)
|
||||
self.assertIsNotNone(head.person_state_adapter[-1].weight.grad)
|
||||
self.assertIsNotNone(head.person_state_adapter_gate.grad)
|
||||
|
||||
config = {
|
||||
"hidden_size": 8,
|
||||
"local_visual_dim": 6,
|
||||
"local_audio_dim": 5,
|
||||
"local_projection_dim": 3,
|
||||
"local_fusion": "gated-person-state-adapter",
|
||||
"person_attention_dim": 4,
|
||||
"person_attention_heads": 2,
|
||||
"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)):
|
||||
@@ -717,7 +850,16 @@ 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", "frame-set", "ego4d"):
|
||||
for mode in (
|
||||
"baseline",
|
||||
"audio",
|
||||
"visual",
|
||||
"av",
|
||||
"frame-set",
|
||||
"f5-screen",
|
||||
"f6-screen",
|
||||
"ego4d",
|
||||
):
|
||||
checkpoint = root / f"{mode}.pth"
|
||||
train = [sys.executable, str(repo / "train_turn_head.py"), *common,
|
||||
"--epochs", "1", "--warmup-epochs", "0", "--batch-size", "4",
|
||||
@@ -734,14 +876,23 @@ class RoundTripTests(unittest.TestCase):
|
||||
train += ["--dev-sessions", "5", "--local-projection-dim", "0"]
|
||||
if mode == "av":
|
||||
train += ["--transition-objective", "r1-supervised"]
|
||||
if mode == "frame-set":
|
||||
person_mode = {
|
||||
"frame-set": "frame-set-attention",
|
||||
"f5-screen": "audio-person-frame-set",
|
||||
"f6-screen": "gated-person-state-adapter",
|
||||
}.get(mode)
|
||||
if person_mode is not None:
|
||||
train += [
|
||||
"--local-fusion", "frame-set-attention",
|
||||
"--local-fusion", person_mode,
|
||||
"--person-attention-dim", "4",
|
||||
"--person-attention-heads", "2",
|
||||
]
|
||||
if mode != "baseline":
|
||||
modality = "av" if mode in ("ego4d", "frame-set") else mode
|
||||
modality = (
|
||||
"av"
|
||||
if mode in ("ego4d", "frame-set", "f5-screen", "f6-screen")
|
||||
else mode
|
||||
)
|
||||
train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality]
|
||||
infer += ["--local-feature-root", str(root / "local")]
|
||||
for command in (train, infer):
|
||||
@@ -752,8 +903,12 @@ 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")
|
||||
if person_mode is not None:
|
||||
self.assertEqual(metrics["local_fusion"], person_mode)
|
||||
self.assertEqual(
|
||||
metrics["local_fusion_scope"],
|
||||
LOCAL_FUSION_SCOPES[person_mode],
|
||||
)
|
||||
# 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