diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index edca672..2c3b91c 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -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 diff --git a/docs/local_feature_baseline.md b/docs/local_feature_baseline.md index b417d17..777c60f 100644 --- a/docs/local_feature_baseline.md +++ b/docs/local_feature_baseline.md @@ -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 批量提取、 diff --git a/docs/progressive_feature_fusion_route_2026-09-02.md b/docs/progressive_feature_fusion_route_2026-09-02.md index 63b72c3..a168721 100644 --- a/docs/progressive_feature_fusion_route_2026-09-02.md +++ b/docs/progressive_feature_fusion_route_2026-09-02.md @@ -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，隔离融合能力。最终再替换为在线人脸或人物 diff --git a/easycom_head.py b/easycom_head.py index bdeff9e..298ea94 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -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,7 +219,33 @@ class TurnPredictionHead(nn.Module): torch.zeros(1, 1, person_attention_dim) ) nn.init.normal_(self.null_person, std=0.02) - fused_size += person_attention_dim + 1 + 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), nn.Linear(fused_size, hidden_size // 2), @@ -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,19 +379,29 @@ 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, ) - states = torch.cat( - (states, audio, audio_valid, person_context), - dim=-1, - ) + 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, person_valid), + dim=-1, + ) elif self.local_input_size: if local is None or local.shape != (*states.shape[:-1], self.local_input_size): raise ValueError( @@ -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, diff --git a/infer_turn_head.py b/infer_turn_head.py index 05e2e19..068c8c8 100644 --- a/infer_turn_head.py +++ b/infer_turn_head.py @@ -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( diff --git a/run_easycom_person_fusion_screening.sh b/run_easycom_person_fusion_screening.sh new file mode 100755 index 0000000..a959d6d --- /dev/null +++ b/run_easycom_person_fusion_screening.sh @@ -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 diff --git a/test_local_turn_features.py b/test_local_turn_features.py index 7334243..03e7103 100644 --- a/test_local_turn_features.py +++ b/test_local_turn_features.py @@ -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)], diff --git a/train_turn_head.py b/train_turn_head.py index a57a80b..b09f6a1 100644 --- a/train_turn_head.py +++ b/train_turn_head.py @@ -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)