feat: add person fusion screening variants

This commit is contained in:
Xin Wang
2026-09-06 21:48:32 +08:00
parent 8e70281330
commit b91a24f308
8 changed files with 447 additions and 47 deletions
+5
View File
@@ -49,6 +49,11 @@ result.
one query over the latest half-open causal person set, retains direct A in
the head, and dynamically pads people per batch. Do not report it as a
measured method until the L20 smoke and full controls pass.
- Two current-sidecar screening controls: `audio-person-frame-set` adds a
zero-gated audio-to-person read before the F1 query, and
`gated-person-state-adapter` adds a zero-gated person residual after the
frozen Thinker readout. They are explicitly recorded as F5/F6 screening,
not full history-conditioned F5 or an internal Thinker F6 adapter.
### Current blocker
+34
View File
@@ -195,6 +195,40 @@ F1 用融合后的 G 和直接局部音频 A 构造一个 query,读取当前
正式比较需要重新训练同 split、seed 和优化设置的 pooled AV 对照,不能直接把不同开发
协议的旧 checkpoint 当作唯一基线。
## 5. F5/F6 screening variants
当前 local sidecar 没有 F3/F4 所需的逐 track 人物历史,预计算的 `states` 也已经位于
Thinker readout 之后。因此下面两项是用于快速筛选的旁路实验,不是路线文档中完整的
F5/F6:
- `audio-person-frame-set`:在 F1 前增加一个 audio-only query 对当前人物集合的读取,
通过零初始化标量 gate 加到 G+A floor query;初始函数与 F1 相同。这是 F5 的
current-frame screening,不含 person-conditioned 5 Hz audio/Omni history。
- `gated-person-state-adapter`:F1 先选择当前人物证据,再通过零初始化标量 gate 将
residual 加到最终 frozen Thinker state。它位于 Thinker **之后**,只能叫 F6 screening,
不能叫 Thinker internal adapter 或 end-to-end Omni adaptation。
两项都继续使用当前 `persons`、audio、5 Hz cutoff 和 NULL-person 合约,无需重新提取
局部特征。训练会逐 epoch 打印 `person_gates`,checkpoint/metrics 会记录明确的
`local_fusion_scope`。运行入口:
```bash
export FEATURE_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack
export TARGET_ROOT=/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame
export LOCAL_ROOT=/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576
# F1、F5 screening、F6 screening;每项 seed 0/1/2,前台顺序执行。
CUDA_VISIBLE_DEVICES=0 bash run_easycom_person_fusion_screening.sh all 0 1 2
# 也可把互不重复的 variant/seed 分配到不同 GPU。
SUMMARIZE=0 CUDA_VISIBLE_DEVICES=1 bash run_easycom_person_fusion_screening.sh f5-screen 0
```
脚本固定 Session 4--11 训练、Session 12 开发选 checkpoint,并输出 Session 1--3 的
exploratory dense mAP 与 `pred/valid`。事件协议冻结前,不能按现有 timing proxy 选择这些
模块;完整 F5 仍需 causal person history,完整 F6 仍需进入 Thinker 上层并与 LoRA 控制
比较。
## Ego4D 接口
完整命令见 [Ego4D 四组局部增强实验](ego4d_local_baselines.md):新增 manifest 批量提取、
@@ -22,8 +22,10 @@
才实施。
实现状态(2026-09-02):F1 的 loader、head、train/infer 路径和回归测试已经加入代码,
但尚未在远程 PyTorch/L20 环境完成 runtime smoke,不能写成已验证结果。F2--F6 和 D1
仍然只是计划。
但尚未在远程 PyTorch/L20 环境完成 runtime smoke,不能写成已验证结果。另加入两个仅用
当前 sidecar 的 F5/F6 **screening** 模式:`audio-person-frame-set` 和
`gated-person-state-adapter`。前者没有 F4 人物历史,后者位于 Thinker readout 之后;二者
都不能写成完整 F5/F6。F2--F4、真正的 F5/F6 和 D1 仍然只是计划。
## 2. 不可改变的实验条件
@@ -214,6 +216,10 @@ G+A floor query 读取。
F5 可增加一次 audio query 到人物 memory 的反向读取,形成单块双向 interaction。必须和
参数匹配 MLP、单向 F4、modality dropout 做比较。若收益不稳定,保留 F4。
代码中的 `audio-person-frame-set` 只用当前帧人物集合和当前 pooled audio 做上述反向读取,
并以零初始化 gate 从 F1 开始优化。它用于判断是否值得继续构建 F3/F4 历史,不满足本节
完整 F5 的输入和对照要求。
F6 才把验证过的人物 memory 通过零初始化 gate 接入 Thinker 上层,并比较:
- 无人物 memory 的参数匹配 Thinker LoRA;
@@ -222,6 +228,10 @@ F6 才把验证过的人物 memory 通过零初始化 gate 接入 Thinker 上层
若 F6 不优于旁路,不为追求端到端表述牺牲可解释性,论文采用更简单的旁路模型。
代码中的 `gated-person-state-adapter` 把当前人物 context 以零初始化 residual 加到已经保存
的最终 Thinker state,只是 post-Thinker screening control。真正 F6 必须重新运行 Omni、
在指定上层 transformer block 内注入 memory,并和参数匹配 LoRA 比较。
## 10. D1:实际在线检测与追踪
方法验证阶段先使用数据集提供的 bbox/track,隔离融合能力。最终再替换为在线人脸或人物
+101 -27
View File
@@ -61,7 +61,18 @@ LAYER_FUSION_MODES = (
"scalar-mixture",
"concat-mlp",
)
LOCAL_FUSION_MODES = ("pooled", "frame-set-attention")
PERSON_SET_FUSION_MODES = (
"frame-set-attention",
"audio-person-frame-set",
"gated-person-state-adapter",
)
LOCAL_FUSION_MODES = ("pooled", *PERSON_SET_FUSION_MODES)
LOCAL_FUSION_SCOPES = {
"pooled": "pooled-baseline",
"frame-set-attention": "f1-current-frame-person-set",
"audio-person-frame-set": "f5-screening-current-frame-no-person-history",
"gated-person-state-adapter": "f6-screening-post-thinker-not-internal",
}
DatasetName = Literal["easycom", "ego4d"]
@@ -132,11 +143,11 @@ class TurnPredictionHead(nn.Module):
raise ValueError("Person attention dimension/heads must be positive")
if person_attention_dim % person_attention_heads:
raise ValueError("Person attention dimension must be divisible by heads")
if local_fusion == "frame-set-attention" and (
if local_fusion in PERSON_SET_FUSION_MODES and (
local_visual_dim <= 0 or local_audio_dim <= 0
):
raise ValueError(
"Frame-set Attention requires both local visual people and audio"
"Person-set fusion requires both local visual people and audio"
)
self.local_fusion = local_fusion
self.person_attention_dim = person_attention_dim
@@ -160,7 +171,7 @@ class TurnPredictionHead(nn.Module):
fused_size = hidden_size
projected_modalities = (
(("audio", local_audio_dim),)
if local_fusion == "frame-set-attention"
if local_fusion in PERSON_SET_FUSION_MODES
else (("visual", local_visual_dim), ("audio", local_audio_dim))
)
for name, dimension in projected_modalities:
@@ -180,7 +191,11 @@ class TurnPredictionHead(nn.Module):
self.person_audio_query: nn.Module | None = None
self.person_attention: nn.MultiheadAttention | None = None
self.null_person: nn.Parameter | None = None
if local_fusion == "frame-set-attention":
self.audio_person_attention: nn.MultiheadAttention | None = None
self.audio_person_gate: nn.Parameter | None = None
self.person_state_adapter: nn.Module | None = None
self.person_state_adapter_gate: nn.Parameter | None = None
if local_fusion in PERSON_SET_FUSION_MODES:
self.person_input_projection = nn.Sequential(
nn.LayerNorm(local_visual_dim),
nn.Linear(local_visual_dim, person_attention_dim),
@@ -204,6 +219,32 @@ class TurnPredictionHead(nn.Module):
torch.zeros(1, 1, person_attention_dim)
)
nn.init.normal_(self.null_person, std=0.02)
if local_fusion == "audio-person-frame-set":
# Screening version of the F5 reverse read: an audio-only
# query first reads the current person memory, then a
# zero-initialized residual gate adds that evidence to the
# regular G+A floor query. This remains a current-frame
# module; it is not the full history-conditioned F5 block.
self.audio_person_attention = nn.MultiheadAttention(
person_attention_dim,
person_attention_heads,
dropout=dropout,
batch_first=True,
)
self.audio_person_gate = nn.Parameter(torch.zeros(()))
if local_fusion == "gated-person-state-adapter":
# Screening version of F6 for precomputed frozen states. It
# injects selected person evidence after the Thinker readout,
# not into a Thinker transformer block. The scalar gate makes
# the visual residual exactly zero at initialization.
self.person_state_adapter = nn.Sequential(
nn.LayerNorm(person_attention_dim),
nn.Linear(person_attention_dim, person_attention_dim),
nn.GELU(),
nn.Linear(person_attention_dim, hidden_size),
)
self.person_state_adapter_gate = nn.Parameter(torch.zeros(()))
else:
fused_size += person_attention_dim + 1
self.net = nn.Sequential(
nn.LayerNorm(fused_size),
@@ -250,7 +291,7 @@ class TurnPredictionHead(nn.Module):
@property
def local_input_size(self) -> int:
if self.local_fusion == "frame-set-attention":
if self.local_fusion in PERSON_SET_FUSION_MODES:
return self.local_audio_dim + 1
return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)
@@ -260,12 +301,12 @@ class TurnPredictionHead(nn.Module):
local: torch.Tensor,
person_features: torch.Tensor | None,
person_mask: torch.Tensor | None,
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
if states.ndim != 2:
raise ValueError("Frame-set Attention expects batched [B,D] states")
raise ValueError("Person-set fusion expects batched [B,D] states")
if local.shape != (len(states), self.local_audio_dim + 1):
raise ValueError(
"Frame-set Attention requires aligned [audio,valid] local features"
"Person-set fusion requires aligned [audio,valid] local features"
)
if (
person_features is None
@@ -274,14 +315,14 @@ class TurnPredictionHead(nn.Module):
or person_features.shape[2] != self.local_visual_dim
):
raise ValueError(
"Frame-set Attention requires person features [B,P,visual_dim]"
"Person-set fusion requires person features [B,P,visual_dim]"
)
if (
person_mask is None
or person_mask.dtype != torch.bool
or person_mask.shape != person_features.shape[:2]
):
raise ValueError("Frame-set Attention requires a boolean person mask [B,P]")
raise ValueError("Person-set fusion requires a boolean person mask [B,P]")
if not torch.isfinite(person_features).all():
raise ValueError("Person features contain NaN/Inf")
assert self.person_input_projection is not None
@@ -293,10 +334,8 @@ class TurnPredictionHead(nn.Module):
audio_values = local[:, : self.local_audio_dim]
audio_valid = local[:, self.local_audio_dim : self.local_audio_dim + 1]
audio_projected = self.local_projections["audio"](audio_values) * audio_valid
query = (
self.person_state_query(states)
+ self.person_audio_query(audio_values) * audio_valid
)
audio_query = self.person_audio_query(audio_values) * audio_valid
query = self.person_state_query(states) + audio_query
people = self.person_input_projection(person_features)
null = self.null_person.expand(len(states), -1, -1)
memory = torch.cat((null, people), dim=1)
@@ -306,6 +345,20 @@ class TurnPredictionHead(nn.Module):
device=person_mask.device,
)
memory_valid = torch.cat((null_mask, person_mask), dim=1)
if self.audio_person_attention is not None:
assert self.audio_person_gate is not None
audio_context, _ = self.audio_person_attention(
audio_query[:, None, :],
memory,
memory,
key_padding_mask=~memory_valid,
need_weights=False,
)
query = query + (
torch.tanh(self.audio_person_gate)
* audio_context[:, 0, :]
* audio_valid
)
context, _ = self.person_attention(
query[:, None, :],
memory,
@@ -314,10 +367,7 @@ class TurnPredictionHead(nn.Module):
need_weights=False,
)
person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype)
return audio_projected, audio_valid, torch.cat(
(context[:, 0, :], person_valid),
dim=-1,
)
return audio_projected, audio_valid, context[:, 0, :], person_valid
def _shared_features(
self,
@@ -329,17 +379,27 @@ class TurnPredictionHead(nn.Module):
"""Return the activity MLP representation shared by both readouts."""
states = self._fuse_layers(states)
if self.local_fusion == "frame-set-attention":
if self.local_fusion in PERSON_SET_FUSION_MODES:
if local is None:
raise ValueError("Frame-set Attention requires local audio features")
audio, audio_valid, person_context = self._frame_set_context(
raise ValueError("Person-set fusion requires local audio features")
audio, audio_valid, person_context, person_valid = self._frame_set_context(
states,
local,
person_features,
person_mask,
)
if self.local_fusion == "gated-person-state-adapter":
assert self.person_state_adapter is not None
assert self.person_state_adapter_gate is not None
states = states + (
torch.tanh(self.person_state_adapter_gate)
* self.person_state_adapter(person_context)
* person_valid
)
states = torch.cat((states, audio, audio_valid), dim=-1)
else:
states = torch.cat(
(states, audio, audio_valid, person_context),
(states, audio, audio_valid, person_context, person_valid),
dim=-1,
)
elif self.local_input_size:
@@ -362,7 +422,7 @@ class TurnPredictionHead(nn.Module):
states = torch.cat(pieces, dim=-1)
elif local is not None and local.shape[-1] != 0:
raise ValueError("Local features supplied to a baseline head")
if self.local_fusion != "frame-set-attention" and (
if self.local_fusion not in PERSON_SET_FUSION_MODES and (
person_features is not None or person_mask is not None
):
raise ValueError("Person sets supplied to a non-frame-set head")
@@ -420,6 +480,20 @@ class TurnPredictionHead(nn.Module):
return None
return self.layer_logits.softmax(dim=0)
def person_interaction_gates(self) -> dict[str, float]:
"""Return bounded residual strengths for screening fusion modes."""
result: dict[str, float] = {}
if self.audio_person_gate is not None:
result["audio_person"] = float(
torch.tanh(self.audio_person_gate.detach()).cpu()
)
if self.person_state_adapter_gate is not None:
result["person_state_adapter"] = float(
torch.tanh(self.person_state_adapter_gate.detach()).cpu()
)
return result
def _validate_loss_inputs(
logits: torch.Tensor,
@@ -1036,10 +1110,10 @@ def load_session_data(
name,
dataset=dataset,
)
if local_fusion == "frame-set-attention":
if local_fusion in PERSON_SET_FUSION_MODES:
if local_modalities != "av":
raise ValueError(
"Frame-set Attention requires --local-modalities av"
"Person-set fusion requires --local-modalities av"
)
from local_turn_features import load_frame_set_features
@@ -1065,7 +1139,7 @@ def load_session_data(
elif local_contract is not None:
raise ValueError("--local-feature-root is required for this checkpoint")
elif local_fusion != "pooled":
raise ValueError("Frame-set Attention requires --local-feature-root")
raise ValueError("Person-set fusion requires --local-feature-root")
return SessionData(
name,
states,
+12
View File
@@ -14,6 +14,7 @@ from tqdm.auto import tqdm
from easycom_head import (
CLASS_NAMES,
DATASET_NAMES,
LOCAL_FUSION_SCOPES,
NUM_CLASSES,
NUM_HORIZONS,
OFFICIAL_TEST_SESSIONS,
@@ -418,6 +419,15 @@ def main() -> None:
print(
f"loss: {loss_name}; output_space: {output_space_for_loss(loss_name)}"
)
print(
f"local fusion: {head.local_fusion}; "
f"scope: {LOCAL_FUSION_SCOPES[head.local_fusion]}"
)
if head.person_interaction_gates():
print(
"person interaction gates: "
+ json.dumps(head.person_interaction_gates(), sort_keys=True)
)
print(f"sessions: {', '.join(sessions)}")
for name in sessions:
data = load_session_data(
@@ -513,6 +523,8 @@ def main() -> None:
"checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
"layer_fusion": head.layer_fusion,
"local_fusion": head.local_fusion,
"local_fusion_scope": LOCAL_FUSION_SCOPES[head.local_fusion],
"person_interaction_gates": head.person_interaction_gates(),
"local_contract": local_contract,
"transition_auxiliary": head.transition_auxiliary,
"transition_objective": checkpoint_payload.get(
+91
View File
@@ -0,0 +1,91 @@
#!/usr/bin/env bash
# Foreground EasyCom F1 plus F5/F6 screening experiments.
# Usage: bash run_easycom_person_fusion_screening.sh [all|f1|f5-screen|f6-screen] [seed ...]
set -euo pipefail
cd "$(dirname "$0")"
FEATURE_ROOT=${FEATURE_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/open_stream_2fps_8s_last4_stack}
TARGET_ROOT=${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/easycom_target_frame}
LOCAL_ROOT=${LOCAL_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_features/local_qwen_1s_encoded_tf4576}
CHECKPOINT_ROOT=${CHECKPOINT_ROOT:-checkpoints/easycom_person_fusion_screening}
RESULT_ROOT=${RESULT_ROOT:-results/easycom_person_fusion_screening}
EPOCHS=${EPOCHS:-4}
WARMUP_EPOCHS=${WARMUP_EPOCHS:-1.6}
LOCAL_PROJECTION_DIM=${LOCAL_PROJECTION_DIM:-256}
PERSON_ATTENTION_DIM=${PERSON_ATTENTION_DIM:-256}
PERSON_ATTENTION_HEADS=${PERSON_ATTENTION_HEADS:-4}
SUMMARIZE=${SUMMARIZE:-1}
export OMP_NUM_THREADS=${OMP_NUM_THREADS:-4}
export MKL_NUM_THREADS=${MKL_NUM_THREADS:-4}
export OPENBLAS_NUM_THREADS=${OPENBLAS_NUM_THREADS:-4}
variant=${1:-all}
if (( $# )); then shift; fi
if (( $# )); then seeds=("$@"); else seeds=(0 1 2); fi
case "$variant" in
all) variants=(f1 f5-screen f6-screen) ;;
f1|f5-screen|f6-screen) variants=("$variant") ;;
*) echo "Expected all, f1, f5-screen, or f6-screen" >&2; exit 2 ;;
esac
mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT"
for variant in "${variants[@]}"; do
case "$variant" in
f1) fusion=frame-set-attention ;;
f5-screen) fusion=audio-person-frame-set ;;
f6-screen) fusion=gated-person-state-adapter ;;
esac
for seed in "${seeds[@]}"; do
checkpoint="$CHECKPOINT_ROOT/${variant}_seed${seed}.pth"
result="$RESULT_ROOT/${variant}_seed${seed}"
if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then
echo "Refusing to overwrite $checkpoint or $result; select new output roots." >&2
exit 1
fi
echo "Training ${variant} (${fusion}), seed=${seed}"
python train_turn_head.py \
--feature-root "$FEATURE_ROOT" \
--target-root "$TARGET_ROOT" \
--local-feature-root "$LOCAL_ROOT" \
--local-modalities av \
--local-fusion "$fusion" \
--local-projection-dim "$LOCAL_PROJECTION_DIM" \
--person-attention-dim "$PERSON_ATTENTION_DIM" \
--person-attention-heads "$PERSON_ATTENTION_HEADS" \
--output "$checkpoint" \
--sessions 4 5 6 7 8 9 10 11 12 \
--dev-sessions 12 \
--epochs "$EPOCHS" \
--batch-size 512 \
--learning-rate 7e-5 \
--weight-decay 5e-5 \
--warmup-epochs "$WARMUP_EPOCHS" \
--warmup-factor 0.3 \
--dropout 0.2 \
--loss lstr-mce \
--transition-objective none \
--layer-fusion concat-mlp \
--expected-mode open-stream-replay \
--seed "$seed" \
--device cuda:0
echo "Inferring ${variant}, seed=${seed}"
python infer_turn_head.py \
--feature-root "$FEATURE_ROOT" \
--target-root "$TARGET_ROOT" \
--local-feature-root "$LOCAL_ROOT" \
--checkpoint "$checkpoint" \
--output-dir "$result" \
--sessions 3 2 1 \
--expected-mode open-stream-replay \
--device cuda:0
done
done
if [[ "$SUMMARIZE" == 1 ]]; then
python summarize_local_baselines.py \
--root "$RESULT_ROOT" \
--dataset easycom \
--protocol causal_valid_only
fi
+162 -7
View File
@@ -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)],
+26 -7
View File
@@ -25,6 +25,7 @@ from easycom_head import (
DATASET_NAMES,
LAYER_FUSION_MODES,
LOCAL_FUSION_MODES,
LOCAL_FUSION_SCOPES,
LOSS_NAMES,
NUM_HORIZONS,
OFFICIAL_TRAIN_SESSIONS,
@@ -64,7 +65,11 @@ def parse_args() -> argparse.Namespace:
"--local-fusion",
choices=LOCAL_FUSION_MODES,
default="pooled",
help="pooled keeps the existing baseline; frame-set-attention queries the latest causal person set.",
help=(
"pooled keeps the existing baseline; frame-set-attention is F1; "
"audio-person-frame-set and gated-person-state-adapter are "
"current-sidecar F5/F6 screening variants, not full F5/F6."
),
)
parser.add_argument(
"--local-projection-dim", type=int, default=256,
@@ -179,12 +184,12 @@ def _validate_args(args: argparse.Namespace) -> None:
raise ValueError("Person attention dimension/heads must be positive")
if args.person_attention_dim % args.person_attention_heads:
raise ValueError("--person-attention-dim must be divisible by --person-attention-heads")
if args.local_fusion == "frame-set-attention" and (
if args.local_fusion != "pooled" and (
args.local_feature_root is None or args.local_modalities != "av"
):
raise ValueError(
"--local-fusion frame-set-attention requires --local-feature-root "
"and --local-modalities av"
"Person-set --local-fusion modes require --local-feature-root and "
"--local-modalities av"
)
if args.epochs <= 0:
raise ValueError("--epochs must be positive")
@@ -505,7 +510,7 @@ def _checkpoint(
history: list[dict[str, Any]],
) -> dict[str, Any]:
return {
"format_version": 5,
"format_version": 6,
"head": head.state_dict(),
"head_config": {
"hidden_size": hidden_size,
@@ -567,6 +572,7 @@ def _checkpoint(
"local_feature_root": str(args.local_feature_root) if args.local_feature_root else None,
"local_modalities": args.local_modalities if args.local_feature_root else None,
"local_fusion": args.local_fusion,
"local_fusion_scope": LOCAL_FUSION_SCOPES[args.local_fusion],
"local_projection_dim": args.local_projection_dim,
"person_attention_dim": args.person_attention_dim,
"person_attention_heads": args.person_attention_heads,
@@ -807,6 +813,10 @@ def main() -> None:
f"fusion_parameters={fusion_parameters:,}, "
f"head_trainable_parameters={trainable_parameters:,}"
)
print(
"local_fusion_scope="
f"{LOCAL_FUSION_SCOPES[args.local_fusion]}"
)
if head.layer_logits is None:
optimizer_parameters: Any = head.parameters()
else:
@@ -967,6 +977,9 @@ def main() -> None:
record["layer_mixture_weights_oldest_to_newest"] = (
layer_weights.detach().float().cpu().tolist()
)
interaction_gates = head.person_interaction_gates()
if interaction_gates:
record["person_interaction_gates"] = interaction_gates
weight_text = (
""
if layer_weights is None
@@ -987,6 +1000,12 @@ def main() -> None:
f"{record['train_expected_transition_reward']:.6f}"
)
)
gate_text = (
""
if not interaction_gates
else ", person_gates="
+ json.dumps(interaction_gates, sort_keys=True)
)
if dev_sessions:
assert dev_states is not None and dev_targets is not None
@@ -1005,7 +1024,7 @@ def main() -> None:
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
f"dev_valid_mean_mAP={dev_map * 100:.3f}"
f"{transition_text}{weight_text}"
f"{transition_text}{weight_text}{gate_text}"
)
if dev_map > float(best_dev_map):
best_dev_map = dev_map
@@ -1031,7 +1050,7 @@ def main() -> None:
else:
print(
f"epoch {epoch:02d}: train_loss={train_loss:.6f}"
f"{transition_text}{weight_text}"
f"{transition_text}{weight_text}{gate_text}"
)
history.append(record)