From 4aed4697a51c6bf849b53c70c90b073388b43df2 Mon Sep 17 00:00:00 2001 From: Xin Wang Date: Mon, 7 Sep 2026 00:30:23 +0800 Subject: [PATCH] Add causal person track GRU residual experiment --- .gitignore | 1 + PROJECT_CONTEXT.md | 28 ++- README.md | 25 +++ docs/README.md | 1 + docs/easycom_track_gru_residual.md | 158 ++++++++++++++ ...ressive_feature_fusion_route_2026-09-02.md | 33 ++- easycom_head.py | 200 ++++++++++++++++- infer_turn_head.py | 15 +- literature_notes/02_egospeak.md | 13 ++ local_turn_features.py | 205 +++++++++++++++++- run_easycom_track_gru_residual.sh | 177 +++++++++++++++ summarize_local_baselines.py | 13 +- test_local_turn_features.py | 159 +++++++++++++- third_party/README.md | 15 ++ train_turn_head.py | 185 +++++++++++++++- 15 files changed, 1185 insertions(+), 43 deletions(-) create mode 100644 docs/easycom_track_gru_residual.md create mode 100755 run_easycom_track_gru_residual.sh create mode 100644 third_party/README.md diff --git a/.gitignore b/.gitignore index 7cf1785..4f06fce 100644 --- a/.gitignore +++ b/.gitignore @@ -11,6 +11,7 @@ __pycache__/ /manifests/ /outputs/ /target_frame/ +/third_party/EgoSpeak/ # Local editors and environments. .idea/ diff --git a/PROJECT_CONTEXT.md b/PROJECT_CONTEXT.md index 2c3b91c..8187baa 100644 --- a/PROJECT_CONTEXT.md +++ b/PROJECT_CONTEXT.md @@ -43,17 +43,27 @@ result. - Ego4D batch and 8-GPU local-feature launchers, with strict frame audits. - Timing-derived wearer floor-state and event diagnostics. +### Measured screening result + +- F1 Frame-set Attention uses fused G+A as one query over the latest half-open + causal person set. On the existing three-seed, four-epoch test screen it + obtained `81.77745 +/- 0.15638` 0.2-second mAP and + `64.95099 +/- 0.09451` ten-horizon mean mAP. This is stable, but is not a + matched improvement over the historical no-dev pooled result; correct/wrong + visual controls remain missing. + ### Implemented, awaiting remote runtime verification -- F1 Frame-set Attention loader/head/train/infer path. It uses fused G+A as - 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. +- `track-gru-residual` builds strictly causal histories from anonymous provided + track IDs, applies one shared GRU per track, lets G+A query current track + states, and adds a zero-gated correction to an exactly initialized pooled + `G+A+V` checkpoint. See + [`docs/easycom_track_gru_residual.md`](docs/easycom_track_gru_residual.md). ### Current blocker @@ -69,12 +79,17 @@ semantic event ground truth. ### Next experiment gates +The user has temporarily chosen to run the dense-AP model screen while P0 is +unfinished. This does not relax the paper claim: event numbers remain proxy +diagnostics until competing-risk censoring is fixed. + 1. **P0:** fix/freeze event attribution, splits and scenario definitions. 2. **P1:** probe aggregate activity and per-person active-speaker grounding; compare correct, wrong-person and time-shifted visual inputs. 3. **P2:** re-evaluate `G`, `G+A`, `G+A+pooled-V` under the frozen event protocol. -4. **P3:** begin the progressive person-fusion route with Frame-set Attention, - then add spatial or track-temporal structure only after its controls pass. +4. **P3:** run the matched pooled/anonymous-track-GRU residual screen, then add + broken-ID, wrong-time and remove-person controls before claiming that track + continuity matters. 5. **P4:** attempt person-memory adapter/Thinker LoRA only if P1/P3 pass. The detailed architecture, controls and stopping gates for P3 are frozen in @@ -227,6 +242,7 @@ Do not reorganize them physically without compatibility wrappers. | `infer_turn_head.py` | Save EgoSpeak-compatible predictions and AP/mAP | | `evaluate_floor_state_changes.py` | CPU floor-state/event proxy evaluation | | `summarize_local_baselines.py` | Aggregate local baseline `metrics.json` files | +| `run_easycom_track_gru_residual.sh` | Matched pooled and anonymous track-GRU training/inference launcher | | `train_omni_turn_lora.py` | Thinker LoRA training from cached AV | | `eval_omni_turn_lora.py` | Thinker LoRA evaluation | diff --git a/README.md b/README.md index b2d2807..29043e2 100644 --- a/README.md +++ b/README.md @@ -1899,6 +1899,31 @@ acceptance rule using Sessions 1--3. LLM-adjudicated candidates are deliberatel kept in a separate JSONL until pilot precision is accepted; they are not yet silently merged into the training arrays. +## Anonymous person-track GRU residual + +The first person-temporal experiment is implemented as +`--local-fusion track-gru-residual`. It preserves an exactly initialized +`G+A+pooled-V` head and learns a zero-gated correction from shared causal GRU +states over temporary person tracks. IDs route history only; they are not model +inputs. Source-clip boundaries, gaps over the configured maximum and unknown +IDs reset history. Existing local Qwen sidecars can be reused without +re-extraction. + +Run the matched pooled/track development protocol with: + +```bash +CUDA_VISIBLE_DEVICES=0 bash run_easycom_track_gru_residual.sh all 0 +``` + +Use separate GPUs for additional seeds, then summarize the common result root +with `summarize_local_baselines.py --paired-deltas`. After method selection is +frozen, `bash run_easycom_track_gru_residual.sh test 0 1 2` evaluates the saved +checkpoints once on Sessions 1--3. The exact causal contract, +EgoSpeak code comparison, L20 smoke command and stopping rule are in +[`docs/easycom_track_gru_residual.md`](docs/easycom_track_gru_residual.md). +This is a frozen-backbone side branch, not an internal Thinker modification or +an EgoSpeak MiniROAD reproduction. + ## Candidate improvement: explicit working and long memory Motivated by EgoSpeak's separation of long-term and working memory, test an diff --git a/docs/README.md b/docs/README.md index 94bcffc..462df5c 100644 --- a/docs/README.md +++ b/docs/README.md @@ -23,6 +23,7 @@ tests. | [`easycom_encoder_only.md`](easycom_encoder_only.md) | Comparing EasyCom Qwen encoder-only A/V/AV representations | | [`ego4d_encoder_only.md`](ego4d_encoder_only.md) | Preparing and running the Ego4D encoder-only comparison | | [`ego4d_local_baselines.md`](ego4d_local_baselines.md) | Previewing Ego4D boxes, extracting local features and running 8-GPU jobs | +| [`easycom_track_gru_residual.md`](easycom_track_gru_residual.md) | Matched pooled-to-track GRU experiment, causal identity contract and L20 commands | ## Design history and implementation options diff --git a/docs/easycom_track_gru_residual.md b/docs/easycom_track_gru_residual.md new file mode 100644 index 0000000..a8b9ece --- /dev/null +++ b/docs/easycom_track_gru_residual.md @@ -0,0 +1,158 @@ +# EasyCom anonymous track-GRU residual + +## Decision + +The first track-temporal experiment keeps the strongest pooled local baseline +as an exact, optionally frozen base and learns only a causal person-track +correction: + +```text +pooled base: G_t + A_t + pooled-V_t -> concat-MLP -> base logits + +per visible person i: + past crop observations on the same temporary track + -> shared projection -> shared one-layer GRU -> H_t,i + +G_t + A_t -> masked attention over {NULL, H_t,1, ..., H_t,N} + -> zero-gated residual logits + +final logits = base logits + tanh(alpha) * person_valid * residual logits +alpha starts at exactly zero +``` + +This is the smallest experiment that asks whether anonymous visual continuity +adds information beyond `G+A+pooled-V`. It is not yet person-conditioned audio +grounding and does not modify the Omni Thinker. + +## What was taken from EgoSpeak + +The official EgoSpeak repository is cloned for inspection at +`third_party/EgoSpeak` and excluded from Git. The inspected revision is +`b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb`. + +Its `MiniROADA` model projects concatenated RGB/audio features, applies a +single-layer GRU, expands the recurrent state into ten anticipation states, +and uses a shared classifier. The EasyCom configuration uses a 128-step +window, stride 4, 2048D embedding, 1024D GRU and ten future steps. + +We reuse the defensible ideas—causal recurrence, shared weights and a ten-step +anticipation output—but not the implementation. Our GRU is shared across +anonymous person tracks, reads only crop observations at or before the current +cutoff, and corrects the already measured pooled Omni baseline. Consequently, +this method must be described as an EgoSpeak-inspired track-temporal adapter, +not an EgoSpeak GRU reproduction. + +## Causal and identity contract + +- A 5 Hz row with cutoff `(t+1)/5` reads the latest local video frame strictly + before that cutoff. +- Track history contains at most `--person-history-frames` real 2 FPS crop + observations and is ordered oldest to newest. +- Temporary IDs only link observations. They are never embedded or passed to + the model, so global ID renaming cannot change a prediction. +- A source-clip boundary always resets a track. A time gap larger than + `--person-track-max-gap-seconds` also resets it. +- Missing, `-1`, `None`, `null` and `unknown` IDs are isolated observations. +- Duplicate known IDs in one frame fail instead of being silently merged. +- Empty person sets produce exactly zero visual correction. + +The existing EasyCom local sidecars already contain crop features, frame times, +source-clip indices and person IDs; this experiment does not require feature +re-extraction. + +## Fair initialization + +`track-gru-residual` requires `--initialize-pooled-checkpoint`. The loader +checks dataset, train/dev split, loss, layer fusion, dimensions and local +feature contract before copying the pooled path. By default those copied +parameters are frozen. The residual gate starts at zero, so before optimization +the track model must reproduce the pooled logits exactly, element for element. + +The matched screening protocol is: + +- EasyCom Session 4--11: training; +- Session 12: development/model selection; +- Session 1--3: final test, run only after history/gap and the method are frozen; +- 4 epochs, three seeds, identical optimizer/loss for pooled and track runs; +- 16 observations at 2 FPS (up to 8 seconds), 1-second maximum link gap. + +This protocol intentionally re-runs the pooled base. Historical no-dev results +cannot be used as the matched control. + +## Run + +After pulling on the L20 machine, run the focused CPU regression tests first: + +```bash +python -m unittest \ + test_local_turn_features.LocalTests.test_track_loader_builds_causal_histories_and_resets_tracks \ + test_local_turn_features.LocalTests.test_track_gru_residual_starts_from_exact_pooled_logits \ + test_local_turn_features.RoundTripTests.test_train_infer_easycom_and_ego4d +``` + +One-seed runtime smoke on an L20: + +```bash +EPOCHS=1 \ +SESSIONS="4 5" \ +DEV_SESSIONS=5 \ +SCREEN_SESSIONS=5 \ +CHECKPOINT_ROOT=checkpoints/easycom_track_gru_smoke \ +RESULT_ROOT=results/easycom_track_gru_smoke \ +CUDA_VISIBLE_DEVICES=0 \ +bash run_easycom_track_gru_residual.sh all 0 +``` + +Then use three GPUs in three foreground terminals. These runs report Session +12 development scores, not the official test set: + +```bash +SUMMARIZE=0 CUDA_VISIBLE_DEVICES=0 bash run_easycom_track_gru_residual.sh all 0 +SUMMARIZE=0 CUDA_VISIBLE_DEVICES=1 bash run_easycom_track_gru_residual.sh all 1 +SUMMARIZE=0 CUDA_VISIBLE_DEVICES=2 bash run_easycom_track_gru_residual.sh all 2 +``` + +After all development jobs finish: + +```bash +python summarize_local_baselines.py \ + --root results/easycom_track_gru_residual \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +``` + +Only after choosing and freezing the method, evaluate the saved checkpoints +once on Sessions 1--3: + +```bash +CUDA_VISIBLE_DEVICES=0 bash run_easycom_track_gru_residual.sh test 0 1 2 +``` + +This writes and summarizes `results/easycom_track_gru_residual_test` without +retraining. A single GPU is sufficient because head inference is inexpensive. + +Paths and hyperparameters can be overridden with the environment variables at +the top of `run_easycom_track_gru_residual.sh`. The launcher runs in the +foreground and refuses to overwrite an existing checkpoint or result. + +## Interpretation and stopping rule + +The initial screen answers only whether track-temporal visual structure can +improve the dense EgoSpeak-compatible metric. Advance this branch only if +same-seed `track_gru - pooled` deltas are directionally consistent on Session +12 without sacrificing short-horizon wearer AP. The official test is a final +confirmation, not another model-selection stage. +Because prior F1 variation was roughly a few tenths of a percentage point, +tiny single-seed changes are not evidence. + +If the track branch wins, next run correct-track versus broken-track, +past-shifted and no-ID temporal controls, followed by person-conditioned audio +grounding. If it does not win, first compare 8/16/32 observations and optionally +unfreeze the pooled base using Session 12 only; do not add deeper cross-attention +merely to rescue a null result. + +The paper-level claim still requires the frozen competing-risk event evaluator +and matched-recall false-entry analysis. Dense AP/mAP improvement alone is a +representation result, not proof that the model understands multi-person floor +behavior. diff --git a/docs/progressive_feature_fusion_route_2026-09-02.md b/docs/progressive_feature_fusion_route_2026-09-02.md index a168721..04b85df 100644 --- a/docs/progressive_feature_fusion_route_2026-09-02.md +++ b/docs/progressive_feature_fusion_route_2026-09-02.md @@ -21,11 +21,14 @@ 融合实验从 F1 Frame-set Attention 开始。后续模块只有在前一级通过预先定义的 gate 后 才实施。 -实现状态（2026-09-02）：F1 的 loader、head、train/infer 路径和回归测试已经加入代码， -但尚未在远程 PyTorch/L20 环境完成 runtime smoke，不能写成已验证结果。另加入两个仅用 +实现状态（更新于 2026-09-07）：F1 已完成三 seed、四 epoch 的远程筛选，得到 +`81.77745 +/- 0.15638` 的 0.2 秒 mAP 和 `64.95099 +/- 0.09451` 的十 horizon mean +mAP。它表现稳定，但与旧 pooled 结果的 split/dev 协议并不匹配，而且尚无 correct/wrong +视觉控制，因此不能写成注意力优于 pooling。另加入两个仅用 当前 sidecar 的 F5/F6 **screening** 模式：`audio-person-frame-set` 和 `gated-person-state-adapter`。前者没有 F4 人物历史，后者位于 Thinker readout 之后；二者 -都不能写成完整 F5/F6。F2--F4、真正的 F5/F6 和 D1 仍然只是计划。 +都不能写成完整 F5/F6。第一版 F3 `track-gru-residual` 已实现但尚待远程运行；F2、F4、 +真正的 F5/F6 和 D1 仍然只是计划。 ## 2. 不可改变的实验条件 @@ -193,6 +196,14 @@ floor-action 结果。 对照包括 frame-set、无 ID temporal pool、正确 track、打断连续性的 track 以及 past-shift。 正确 track 必须优于无 ID 和打断 track，才能说明身份连续性有价值。 +第一版实现为 `track-gru-residual`。它没有直接替换 pooled 头，而是从同 seed、同 split、 +同 local contract 的 `G+A+pooled-V` checkpoint 精确复制并默认冻结 base path，再增加共享 +track GRU、G+A 对 track states 的 masked attention 以及零初始化 scalar gate。这样训练 +开始前 logits 与 pooled baseline 逐元素相同，空人物时 correction 恒为零，实验首先回答 +“track continuity 是否提供增量信息”，而不会把随机重训 base 的波动误当作 GRU 收益。 +实现、契约和命令见 +[`easycom_track_gru_residual.md`](easycom_track_gru_residual.md)。 + ## 8. F4：Person-conditioned audio/Omni history F4 让每个匿名人物分别读取已有的 5 Hz 因果历史： @@ -268,11 +279,11 @@ README 必须写明 split、seed、checkpoint、输入特征版本、评测协 ## 12. 最近执行顺序 -1. 完成并测试 P0 competing-risk censoring，人工打印至少一条 response 和 yield 时间轴。 -2. 扩展 local sidecar loader，使其在不破坏现有 pooled 加载路径的情况下返回逐帧人物 - feature、offset 和 mask；测试成功、失败、空人物和越界路径。 -3. 实现 F1 最小 Frame-set Attention，不加入 bbox 或 track temporal。 -4. 单 session/短片 smoke：检查 shape、cutoff、chosen frame time、mask、空人物和排列 - 不变性。 -5. 在固定 train-internal dev 上运行 B1、B3、F1-a、F1-c、F1-d、F1-e 和 F1-g。 -6. F1 通过后才进入 F2；若未通过，优先审计人物 crop 和 AV 对齐，不继续堆模型。 +1. 用户当前选择先做 dense-AP model screen；运行 matched pooled 与 + `track-gru-residual`，固定 Session 12 dev，不用测试集选历史长度。 +2. 若 track GRU 有稳定增益，补 broken-track、past-shift、no-ID 和 remove-person 控制；若 + 无增益，只在 dev 上比较 8/16/32 observations 后停止该分支。 +3. correct track 通过控制后，再实现 F4 person-conditioned audio/Omni history；否则回到 + 对齐审计，不直接堆更深 interaction。 +4. 在形成论文事件结论前，仍须完成 P0 competing-risk censoring，并人工打印 response 与 + yield 时间轴。当前事件结果只能称为 timing-derived proxy。 diff --git a/easycom_head.py b/easycom_head.py index 298ea94..32f0a88 100644 --- a/easycom_head.py +++ b/easycom_head.py @@ -61,17 +61,23 @@ LAYER_FUSION_MODES = ( "scalar-mixture", "concat-mlp", ) -PERSON_SET_FUSION_MODES = ( +CURRENT_FRAME_PERSON_SET_FUSION_MODES = ( "frame-set-attention", "audio-person-frame-set", "gated-person-state-adapter", ) +TRACK_PERSON_FUSION_MODES = ("track-gru-residual",) +PERSON_SET_FUSION_MODES = ( + *CURRENT_FRAME_PERSON_SET_FUSION_MODES, + *TRACK_PERSON_FUSION_MODES, +) 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", + "track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual", } DatasetName = Literal["easycom", "ego4d"] @@ -95,6 +101,8 @@ class TurnPredictionHead(nn.Module): local_fusion: str = "pooled", person_attention_dim: int = 256, person_attention_heads: int = 4, + person_history_frames: int = 16, + person_track_max_gap_seconds: float = 1.0, ) -> None: super().__init__() if num_input_layers <= 0: @@ -143,6 +151,10 @@ 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 person_history_frames <= 0: + raise ValueError("person_history_frames must be positive") + if person_track_max_gap_seconds <= 0: + raise ValueError("person_track_max_gap_seconds must be positive") if local_fusion in PERSON_SET_FUSION_MODES and ( local_visual_dim <= 0 or local_audio_dim <= 0 ): @@ -152,6 +164,8 @@ class TurnPredictionHead(nn.Module): self.local_fusion = local_fusion self.person_attention_dim = person_attention_dim self.person_attention_heads = person_attention_heads + self.person_history_frames = person_history_frames + self.person_track_max_gap_seconds = person_track_max_gap_seconds if self.learnable_layer_mixture: # Zero initialization exactly reproduces a uniform layer mean at # step zero while allowing training to select useful query layers. @@ -171,7 +185,7 @@ class TurnPredictionHead(nn.Module): fused_size = hidden_size projected_modalities = ( (("audio", local_audio_dim),) - if local_fusion in PERSON_SET_FUSION_MODES + if local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES else (("visual", local_visual_dim), ("audio", local_audio_dim)) ) for name, dimension in projected_modalities: @@ -195,6 +209,9 @@ class TurnPredictionHead(nn.Module): 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 + self.person_track_gru: nn.GRU | None = None + self.person_track_delta: nn.Module | None = None + self.person_track_residual_gate: nn.Parameter | None = None if local_fusion in PERSON_SET_FUSION_MODES: self.person_input_projection = nn.Sequential( nn.LayerNorm(local_visual_dim), @@ -244,6 +261,25 @@ class TurnPredictionHead(nn.Module): nn.Linear(person_attention_dim, hidden_size), ) self.person_state_adapter_gate = nn.Parameter(torch.zeros(())) + elif local_fusion == "track-gru-residual": + # Apply one shared causal GRU independently to every anonymous + # face track. The regular pooled G+A+V head remains the exact + # base path; this branch can only add a zero-gated correction. + self.person_track_gru = nn.GRU( + person_attention_dim, + person_attention_dim, + num_layers=1, + batch_first=True, + ) + delta_input = hidden_size // 2 + person_attention_dim + 1 + self.person_track_delta = nn.Sequential( + nn.LayerNorm(delta_input), + nn.Linear(delta_input, hidden_size // 2), + nn.GELU(), + nn.Dropout(dropout), + nn.Linear(hidden_size // 2, num_horizons * num_classes), + ) + self.person_track_residual_gate = nn.Parameter(torch.zeros(())) else: fused_size += person_attention_dim + 1 self.net = nn.Sequential( @@ -291,7 +327,7 @@ class TurnPredictionHead(nn.Module): @property def local_input_size(self) -> int: - if self.local_fusion in PERSON_SET_FUSION_MODES: + if self.local_fusion in CURRENT_FRAME_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) @@ -369,6 +405,95 @@ class TurnPredictionHead(nn.Module): person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype) return audio_projected, audio_valid, context[:, 0, :], person_valid + def _track_gru_context( + self, + states: torch.Tensor, + local: torch.Tensor, + person_histories: torch.Tensor | None, + history_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """Encode each causal track with a shared GRU, then query the set.""" + + if states.ndim != 2: + raise ValueError("Track GRU fusion expects batched [B,D] states") + if local.shape != (len(states), self.local_input_size): + raise ValueError("Track GRU fusion requires pooled [visual,audio,valid] input") + if ( + person_histories is None + or person_histories.ndim != 4 + or person_histories.shape[0] != len(states) + or person_histories.shape[2] != self.person_history_frames + or person_histories.shape[3] != self.local_visual_dim + ): + raise ValueError( + "Track GRU fusion requires person histories " + "[B,P,person_history_frames,visual_dim]" + ) + if ( + history_mask is None + or history_mask.dtype != torch.bool + or history_mask.shape != person_histories.shape[:3] + ): + raise ValueError("Track GRU fusion requires a boolean history mask [B,P,L]") + if not torch.isfinite(person_histories).all(): + raise ValueError("Person track histories contain NaN/Inf") + if bool((history_mask[:, :, 1:] & ~history_mask[:, :, :-1]).any()): + raise ValueError("Person track histories must be left aligned") + assert self.person_input_projection is not None + assert self.person_state_query is not None + assert self.person_audio_query is not None + assert self.person_attention is not None + assert self.null_person is not None + assert self.person_track_gru is not None + + batch, people, length, _ = person_histories.shape + person_valid = history_mask.any(dim=-1) + if people: + projected = self.person_input_projection(person_histories) + flat = projected.reshape(batch * people, length, -1) + lengths = history_mask.sum(dim=-1).reshape(-1) + safe_lengths = lengths.clamp_min(1).to(device="cpu") + packed = nn.utils.rnn.pack_padded_sequence( + flat, + safe_lengths, + batch_first=True, + enforce_sorted=False, + ) + _, hidden = self.person_track_gru(packed) + track_states = hidden[-1].reshape(batch, people, -1) + track_states = track_states * person_valid[..., None] + else: + track_states = states.new_zeros( + (batch, 0, self.person_attention_dim) + ) + + visual_width = self.local_visual_dim + 1 + audio_values = local[:, visual_width : visual_width + self.local_audio_dim] + audio_valid = local[:, visual_width + self.local_audio_dim :] + if audio_valid.shape != (batch, 1): + raise AssertionError("Unexpected pooled AV local feature layout") + query = self.person_state_query(states) + query = query + self.person_audio_query(audio_values) * audio_valid + memory = torch.cat( + (self.null_person.expand(batch, -1, -1), track_states), dim=1 + ) + memory_valid = torch.cat( + ( + torch.ones((batch, 1), dtype=torch.bool, device=states.device), + person_valid, + ), + dim=1, + ) + context, _ = self.person_attention( + query[:, None, :], + memory, + memory, + key_padding_mask=~memory_valid, + need_weights=False, + ) + any_person = person_valid.any(dim=1, keepdim=True).to(states.dtype) + return context[:, 0, :] * any_person, any_person + def _shared_features( self, states: torch.Tensor, @@ -379,7 +504,7 @@ class TurnPredictionHead(nn.Module): """Return the activity MLP representation shared by both readouts.""" states = self._fuse_layers(states) - if self.local_fusion in PERSON_SET_FUSION_MODES: + if self.local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES: if local is None: raise ValueError("Person-set fusion requires local audio features") audio, audio_valid, person_context, person_valid = self._frame_set_context( @@ -428,6 +553,42 @@ class TurnPredictionHead(nn.Module): raise ValueError("Person sets supplied to a non-frame-set head") return self.net[:-1](states) + def _activity_logits( + self, + states: torch.Tensor, + local: torch.Tensor | None, + person_features: torch.Tensor | None, + person_mask: torch.Tensor | None, + ) -> tuple[torch.Tensor, torch.Tensor]: + features = self._shared_features( + states, + local, + person_features, + person_mask, + ) + flat_logits = self.net[-1](features) + if self.local_fusion == "track-gru-residual": + if local is None: + raise ValueError("Track GRU residual requires pooled local AV") + assert self.person_track_delta is not None + assert self.person_track_residual_gate is not None + fused_states = self._fuse_layers(states) + context, person_valid = self._track_gru_context( + fused_states, + local, + person_features, + person_mask, + ) + delta = self.person_track_delta( + torch.cat((features, context, person_valid), dim=-1) + ) + flat_logits = flat_logits + ( + torch.tanh(self.person_track_residual_gate) + * delta + * person_valid + ) + return features, flat_logits + def forward( self, states: torch.Tensor, @@ -435,13 +596,12 @@ class TurnPredictionHead(nn.Module): person_features: torch.Tensor | None = None, person_mask: torch.Tensor | None = None, ) -> torch.Tensor: - features = self._shared_features( + features, logits = self._activity_logits( states, local, person_features, person_mask, ) - logits = self.net[-1](features) return logits.view( *features.shape[:-1], self.num_horizons, self.num_classes ) @@ -459,13 +619,13 @@ class TurnPredictionHead(nn.Module): raise RuntimeError( "This head was created without transition_auxiliary=True" ) - features = self._shared_features( + features, flat_activity_logits = self._activity_logits( states, local, person_features, person_mask, ) - activity_logits = self.net[-1](features).view( + activity_logits = flat_activity_logits.view( *features.shape[:-1], self.num_horizons, self.num_classes, @@ -492,6 +652,10 @@ class TurnPredictionHead(nn.Module): result["person_state_adapter"] = float( torch.tanh(self.person_state_adapter_gate.detach()).cpu() ) + if self.person_track_residual_gate is not None: + result["person_track_residual"] = float( + torch.tanh(self.person_track_residual_gate.detach()).cpu() + ) return result @@ -1004,6 +1168,7 @@ def load_session_data( local_modalities: str = "av", local_contract: dict[str, Any] | None = None, local_fusion: str = "pooled", + person_track_max_gap_seconds: float = 1.0, ) -> SessionData: """Load and validate one saved Omni feature payload and target array.""" @@ -1110,7 +1275,22 @@ def load_session_data( name, dataset=dataset, ) - if local_fusion in PERSON_SET_FUSION_MODES: + if local_fusion in TRACK_PERSON_FUSION_MODES: + if local_modalities != "av": + raise ValueError( + "Track person fusion requires --local-modalities av" + ) + from local_turn_features import load_track_features + + local, local_people, contract = load_track_features( + local_path, + payload, + sample_id=name, + dataset=dataset, + track_max_gap_seconds=person_track_max_gap_seconds, + expected_contract=local_contract, + ) + elif local_fusion in CURRENT_FRAME_PERSON_SET_FUSION_MODES: if local_modalities != "av": raise ValueError( "Person-set fusion requires --local-modalities av" @@ -1196,6 +1376,8 @@ def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]: config.setdefault("local_fusion", "pooled") config.setdefault("person_attention_dim", 256) config.setdefault("person_attention_heads", 4) + config.setdefault("person_history_frames", 16) + config.setdefault("person_track_max_gap_seconds", 1.0) config.setdefault( "layer_fusion", ( diff --git a/infer_turn_head.py b/infer_turn_head.py index 068c8c8..7f0ba01 100644 --- a/infer_turn_head.py +++ b/infer_turn_head.py @@ -149,7 +149,13 @@ def _predict_states( start, min(start + batch_size, len(states)), ) - person_batch, person_mask = local_person_features.padded(rows) + if head.local_fusion == "track-gru-residual": + person_batch, person_mask = local_person_features.padded_history( + rows, + max_observations=head.person_history_frames, + ) + else: + person_batch, person_mask = local_person_features.padded(rows) person_batch = person_batch.to( device=device, dtype=torch.float32, @@ -440,6 +446,7 @@ def main() -> None: local_modalities=local_contract["modalities"] if local_contract else "av", local_contract=local_contract, local_fusion=head.local_fusion, + person_track_max_gap_seconds=head.person_track_max_gap_seconds, ) if data.hidden_size != int(head_config["hidden_size"]): raise ValueError( @@ -519,12 +526,18 @@ def main() -> None: "prediction_layout": "[target_time, class, future_horizon]", "checkpoint": str(args.checkpoint.resolve()), "checkpoint_epoch": checkpoint_payload.get("epoch"), + "checkpoint_training_epochs": checkpoint_payload.get( + "training_config", + {}, + ).get("epochs"), "checkpoint_train_sessions": checkpoint_payload.get("train_sessions"), "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(), + "person_history_frames": head.person_history_frames, + "person_track_max_gap_seconds": head.person_track_max_gap_seconds, "local_contract": local_contract, "transition_auxiliary": head.transition_auxiliary, "transition_objective": checkpoint_payload.get( diff --git a/literature_notes/02_egospeak.md b/literature_notes/02_egospeak.md index a011003..b6d0e25 100644 --- a/literature_notes/02_egospeak.md +++ b/literature_notes/02_egospeak.md @@ -3,6 +3,8 @@ - 会议：Findings of NAACL 2025 - 阅读深度：深入阅读 - 官方链接：https://aclanthology.org/2025.findings-naacl.163/ +- 官方代码：https://github.com/jun297/EgoSpeak +- 本项目核对的代码 revision：`b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb` ## 摘要（中文转述） @@ -12,3 +14,14 @@ EgoSpeak 研究第一人称智能体在自然多人对话中“何时开口” 它定义了本项目复现的核心协议：5 Hz 目标、多个未来预测时距、EasyCom/Ego4D 划分及 AP/mAP 评测。我们发现 EasyCom 中 wearer 与 others 可以同时为真，因此三路 softmax 的互斥输出与真实标签结构不完全一致；四状态 softmax 和独立 BCE 的消融虽更匹配标签结构，但当前四个 epoch 的收益很小，说明主要瓶颈可能不是输出空间。 +## GRU 实现核对 + +官方 `MiniROADA` 将同一时刻的 RGB 与音频 feature 拼接，经 +`Linear -> LayerNorm -> ReLU -> Dropout` 后送入单层 GRU，再把每个 recurrent state +映射成十个 anticipation hidden states 并共享分类器。EasyCom 配置使用 128 个 5 Hz +step（25.6 秒）、stride 4、2048D embedding、1024D hidden 和 10 个 future steps。 + +本项目的 `track-gru-residual` 没有复制这段实现，也不是该基线的复现：它把 GRU 参数在 +匿名人物轨迹之间共享，并用零门控残差修正 matched pooled Omni head。官方实现仍应作为 +单独的 sequence-GRU 强基线；逐人 track continuity 是否有效，需要 broken-ID/no-ID 控制 +而不能由架构名称推断。 diff --git a/local_turn_features.py b/local_turn_features.py index d6c2f6f..b020592 100644 --- a/local_turn_features.py +++ b/local_turn_features.py @@ -25,6 +25,7 @@ class FramePersonFeatures: frame_offsets: torch.LongTensor row_to_frame: torch.LongTensor chosen_frame_times: torch.Tensor + previous_observation: torch.LongTensor | None = None def __post_init__(self) -> None: if ( @@ -65,6 +66,19 @@ class FramePersonFeatures: raise ValueError("Only rows without an available frame may have NaN frame time") if not torch.isfinite(self.chosen_frame_times[~missing]).all(): raise ValueError("Chosen person frame times must be finite") + if self.previous_observation is not None: + previous = self.previous_observation + if ( + not isinstance(previous, torch.Tensor) + or previous.dtype != torch.long + or previous.shape != (len(self.features),) + ): + raise ValueError("previous_observation must be int64 [N]") + indices = torch.arange(len(previous), dtype=torch.long) + if bool(((previous < -1) | (previous >= indices)).any()): + raise ValueError( + "Each previous observation must be -1 or point strictly backward" + ) def __len__(self) -> int: return len(self.row_to_frame) @@ -89,13 +103,14 @@ class FramePersonFeatures: self.frame_offsets, self.row_to_frame[indices], self.chosen_frame_times[indices], + self.previous_observation, ) - def padded( + def _padded_observation_indices( self, row_indices: torch.Tensor | None = None, - ) -> tuple[torch.Tensor, torch.BoolTensor]: - """Materialize only one batch as [B,P,D] plus a True-for-real mask.""" + ) -> tuple[torch.LongTensor, torch.BoolTensor]: + """Return current-frame observation indices and a real-person mask.""" if row_indices is None: frames = self.row_to_frame @@ -121,12 +136,77 @@ class FramePersonFeatures: width = int(counts.max()) if len(counts) else 0 mask = torch.arange(width)[None, :] < counts[:, None] if width == 0: - return self.features.new_zeros((len(frames), 0, self.dimension)), mask - person_indices = starts[:, None] + torch.arange(width)[None, :] - safe_indices = person_indices.masked_fill(~mask, 0) + return torch.zeros((len(frames), 0), dtype=torch.long), mask + indices = starts[:, None] + torch.arange(width)[None, :] + return indices.masked_fill(~mask, 0), mask + + def padded( + self, + row_indices: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.BoolTensor]: + """Materialize only one batch as [B,P,D] plus a True-for-real mask.""" + + safe_indices, mask = self._padded_observation_indices(row_indices) + if safe_indices.shape[1] == 0: + return self.features.new_zeros((len(safe_indices), 0, self.dimension)), mask values = self.features[safe_indices] return values.masked_fill(~mask[..., None], 0), mask + def padded_history( + self, + row_indices: torch.Tensor | None = None, + *, + max_observations: int, + ) -> tuple[torch.Tensor, torch.BoolTensor]: + """Return causal per-current-person histories as ``[B,P,L,D]``. + + Histories are chronological and left aligned. Track links are built + from provided IDs with an explicit causal gap/reset policy when the + sidecar is loaded; no identity value is exposed to the model. + """ + + if max_observations <= 0: + raise ValueError("max_observations must be positive") + if self.previous_observation is None: + raise ValueError("Track histories were not constructed for this person set") + current, current_mask = self._padded_observation_indices(row_indices) + batch, people = current.shape + if people == 0: + return ( + self.features.new_zeros( + (batch, 0, max_observations, self.dimension) + ), + torch.zeros( + (batch, 0, max_observations), dtype=torch.bool + ), + ) + + reverse_indices = torch.zeros( + (batch, people, max_observations), dtype=torch.long + ) + reverse_mask = torch.zeros( + (batch, people, max_observations), dtype=torch.bool + ) + cursor = current.clone() + cursor_valid = current_mask.clone() + for step in range(max_observations): + reverse_indices[:, :, step] = cursor.masked_fill(~cursor_valid, 0) + reverse_mask[:, :, step] = cursor_valid + safe = cursor.masked_fill(~cursor_valid, 0) + cursor = self.previous_observation[safe] + cursor_valid &= cursor >= 0 + + lengths = reverse_mask.sum(dim=-1) + output_positions = torch.arange(max_observations)[None, None, :] + source_positions = lengths[:, :, None] - 1 - output_positions + history_mask = output_positions < lengths[:, :, None] + source_positions = source_positions.clamp_min(0) + chronological = reverse_indices.gather(dim=-1, index=source_positions) + chronological = chronological.masked_fill(~history_mask, 0) + histories = self.features[chronological] + histories = histories.masked_fill(~history_mask[..., None], 0) + return histories, history_mask + @classmethod def concatenate( cls, @@ -143,6 +223,10 @@ class FramePersonFeatures: time_parts = [] feature_base = 0 frame_base = 0 + previous_parts = [] + has_track_links = [part.previous_observation is not None for part in parts] + if any(has_track_links) and not all(has_track_links): + raise ValueError("Track history links are missing for some sessions") for part in parts: feature_parts.append(part.features) offset_parts.append(part.frame_offsets[1:] + feature_base) @@ -151,6 +235,11 @@ class FramePersonFeatures: rows[valid] += frame_base row_parts.append(rows) time_parts.append(part.chosen_frame_times) + if part.previous_observation is not None: + previous = part.previous_observation.clone() + valid_previous = previous >= 0 + previous[valid_previous] += feature_base + previous_parts.append(previous) feature_base += len(part.features) frame_base += len(part.frame_offsets) - 1 return cls( @@ -158,6 +247,7 @@ class FramePersonFeatures: torch.cat(offset_parts, dim=0), torch.cat(row_parts, dim=0), torch.cat(time_parts, dim=0), + torch.cat(previous_parts, dim=0) if previous_parts else None, ) @@ -257,6 +347,7 @@ def load_frame_set_features( sample_id: str, dataset: str, expected_contract: dict[str, Any] | None = None, + track_max_gap_seconds: float | None = None, ) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]: """Load direct audio and the latest causal per-frame person set per 5 Hz row.""" @@ -340,11 +431,23 @@ def load_frame_set_features( chosen[has_frame] = frame_times[row_to_frame[has_frame]] if bool((chosen[has_frame] >= times[has_frame]).any()): raise ValueError(f"Person frame mapping selected future visual evidence: {path}") + previous_observation = None + if track_max_gap_seconds is not None: + if not math.isfinite(track_max_gap_seconds) or track_max_gap_seconds <= 0: + raise ValueError("track_max_gap_seconds must be finite and positive") + previous_observation = _build_previous_track_links( + records, + frame_times, + frame_offsets, + max_gap_seconds=track_max_gap_seconds, + path=path, + ) frame_sets = FramePersonFeatures( features, frame_offsets, row_to_frame.to(dtype=torch.long), chosen, + previous_observation, ) direct_audio = torch.cat( (audio.float(), audio_valid[:, None].float()), @@ -353,6 +456,96 @@ def load_frame_set_features( return direct_audio, frame_sets, contract +def _build_previous_track_links( + records: list[Any], + frame_times: torch.Tensor, + frame_offsets: torch.LongTensor, + *, + max_gap_seconds: float, + path: Path, +) -> torch.LongTensor: + """Link observations causally by provided track ID without exposing IDs. + + A source-clip boundary always resets a track. Missing/unknown IDs are + isolated observations instead of being incorrectly merged into one track. + Duplicate known IDs within one frame fail closed. + """ + + previous = torch.full((len(records),), -1, dtype=torch.long) + last_by_track: dict[tuple[int, str], tuple[int, float]] = {} + unknown_values = {"", "-1", "none", "null", "unknown"} + for frame_index, time_value in enumerate(frame_times.tolist()): + start = int(frame_offsets[frame_index]) + end = int(frame_offsets[frame_index + 1]) + seen: set[tuple[int, str]] = set() + for observation_index in range(start, end): + record = records[observation_index] + if not isinstance(record, dict): + raise ValueError(f"Invalid person record at {observation_index}: {path}") + raw_id = record.get("id") + text_id = "" if raw_id is None else str(raw_id).strip() + source_clip_index = record.get("source_clip_index", 0) + if isinstance(source_clip_index, bool) or not isinstance( + source_clip_index, int + ): + raise ValueError( + f"Invalid source_clip_index at person record " + f"{observation_index}: {path}" + ) + if text_id.lower() in unknown_values: + # Never join multiple anonymous detections under a shared -1. + continue + key = (source_clip_index, text_id) + if key in seen: + raise ValueError( + f"Duplicate person/track ID {text_id!r} in frame " + f"{frame_index}: {path}" + ) + seen.add(key) + old = last_by_track.get(key) + if old is not None: + old_index, old_time = old + gap = float(time_value) - old_time + if gap <= 0: + raise ValueError(f"Non-causal person track order: {path}") + if gap <= max_gap_seconds + 1e-8: + previous[observation_index] = old_index + last_by_track[key] = (observation_index, float(time_value)) + return previous + + +def load_track_features( + path: Path, + reference: dict[str, Any], + *, + sample_id: str, + dataset: str, + track_max_gap_seconds: float, + expected_contract: dict[str, Any] | None = None, +) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]: + """Load the exact pooled AV baseline plus causal anonymous track links.""" + + pooled, contract = load_local_features( + path, + reference, + sample_id=sample_id, + dataset=dataset, + modalities="av", + expected_contract=expected_contract, + ) + _, people, second_contract = load_frame_set_features( + path, + reference, + sample_id=sample_id, + dataset=dataset, + expected_contract=contract, + track_max_gap_seconds=track_max_gap_seconds, + ) + if second_contract != contract: + raise AssertionError("Track and pooled local contracts unexpectedly differ") + return pooled, people, contract + + def local_head_dimensions(contract: dict[str, Any] | None) -> dict[str, int]: if contract is None: return {"local_visual_dim": 0, "local_audio_dim": 0} diff --git a/run_easycom_track_gru_residual.sh b/run_easycom_track_gru_residual.sh new file mode 100755 index 0000000..703cd43 --- /dev/null +++ b/run_easycom_track_gru_residual.sh @@ -0,0 +1,177 @@ +#!/usr/bin/env bash +# Foreground matched EasyCom pooled AV -> causal anonymous-track GRU residual. +# Usage: bash run_easycom_track_gru_residual.sh [all|pooled|track|test] [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_track_gru_residual} +RESULT_ROOT=${RESULT_ROOT:-results/easycom_track_gru_residual} +TEST_RESULT_ROOT=${TEST_RESULT_ROOT:-results/easycom_track_gru_residual_test} +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} +PERSON_HISTORY_FRAMES=${PERSON_HISTORY_FRAMES:-16} +PERSON_TRACK_MAX_GAP_SECONDS=${PERSON_TRACK_MAX_GAP_SECONDS:-1.0} +SUMMARIZE=${SUMMARIZE:-1} +read -r -a SESSION_ARGS <<< "${SESSIONS:-4 5 6 7 8 9 10 11 12}" +read -r -a DEV_SESSION_ARGS <<< "${DEV_SESSIONS:-12}" +read -r -a SCREEN_SESSION_ARGS <<< "${SCREEN_SESSIONS:-12}" +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=(pooled track) ;; + pooled|track) variants=("$variant") ;; + test) variants=(pooled track) ;; + *) echo "Expected all, pooled, track, or test" >&2; exit 2 ;; +esac + +mkdir -p "$CHECKPOINT_ROOT" "$RESULT_ROOT" + +infer_checkpoint() { + local checkpoint=$1 result=$2 + shift 2 + if [[ -e "$result/metrics.json" ]]; then + echo "Refusing to overwrite $result/metrics.json" >&2 + return 1 + fi + python infer_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --local-feature-root "$LOCAL_ROOT" \ + --checkpoint "$checkpoint" \ + --output-dir "$result" \ + --sessions "$@" \ + --expected-mode open-stream-replay \ + --device cuda:0 +} + +train_and_infer_pooled() { + local seed=$1 checkpoint result + checkpoint="$CHECKPOINT_ROOT/pooled_seed${seed}.pth" + result="$RESULT_ROOT/pooled_seed${seed}" + if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then + echo "Refusing to overwrite $checkpoint or $result" >&2 + return 1 + fi + python train_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --local-feature-root "$LOCAL_ROOT" \ + --local-modalities av \ + --local-fusion pooled \ + --local-projection-dim "$LOCAL_PROJECTION_DIM" \ + --output "$checkpoint" \ + --sessions "${SESSION_ARGS[@]}" \ + --dev-sessions "${DEV_SESSION_ARGS[@]}" \ + --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 + # Screen architecture/hyperparameters only on the held-out training session. + infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}" +} + +train_and_infer_track() { + local seed=$1 base checkpoint result + base="$CHECKPOINT_ROOT/pooled_seed${seed}.pth" + checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth" + result="$RESULT_ROOT/track_gru_seed${seed}" + if [[ ! -s "$base" ]]; then + echo "Missing matched pooled checkpoint: $base; run variant pooled or all first" >&2 + return 1 + fi + if [[ -e "$checkpoint" || -e "$result/metrics.json" ]]; then + echo "Refusing to overwrite $checkpoint or $result" >&2 + return 1 + fi + python train_turn_head.py \ + --feature-root "$FEATURE_ROOT" \ + --target-root "$TARGET_ROOT" \ + --local-feature-root "$LOCAL_ROOT" \ + --local-modalities av \ + --local-fusion track-gru-residual \ + --local-projection-dim "$LOCAL_PROJECTION_DIM" \ + --person-attention-dim "$PERSON_ATTENTION_DIM" \ + --person-attention-heads "$PERSON_ATTENTION_HEADS" \ + --person-history-frames "$PERSON_HISTORY_FRAMES" \ + --person-track-max-gap-seconds "$PERSON_TRACK_MAX_GAP_SECONDS" \ + --initialize-pooled-checkpoint "$base" \ + --output "$checkpoint" \ + --sessions "${SESSION_ARGS[@]}" \ + --dev-sessions "${DEV_SESSION_ARGS[@]}" \ + --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 + infer_checkpoint "$checkpoint" "$result" "${SCREEN_SESSION_ARGS[@]}" +} + +if [[ "$variant" == test ]]; then + mkdir -p "$TEST_RESULT_ROOT" + for selected in "${variants[@]}"; do + for seed in "${seeds[@]}"; do + checkpoint="$CHECKPOINT_ROOT/${selected}_seed${seed}.pth" + if [[ "$selected" == track ]]; then + checkpoint="$CHECKPOINT_ROOT/track_gru_seed${seed}.pth" + result="$TEST_RESULT_ROOT/track_gru_seed${seed}" + else + result="$TEST_RESULT_ROOT/pooled_seed${seed}" + fi + if [[ ! -s "$checkpoint" ]]; then + echo "Missing checkpoint: $checkpoint" >&2 + exit 1 + fi + echo "$(date): final test ${selected}, seed=${seed}" + infer_checkpoint "$checkpoint" "$result" 3 2 1 + done + done + summary_root=$TEST_RESULT_ROOT +else + for selected in "${variants[@]}"; do + for seed in "${seeds[@]}"; do + echo "$(date): ${selected}, seed=${seed}" + if [[ "$selected" == pooled ]]; then + train_and_infer_pooled "$seed" + else + train_and_infer_track "$seed" + fi + done + done + summary_root=$RESULT_ROOT +fi + +if [[ "$SUMMARIZE" == 1 ]]; then + python summarize_local_baselines.py \ + --root "$summary_root" \ + --dataset easycom \ + --protocol causal_valid_only \ + --paired-deltas +fi diff --git a/summarize_local_baselines.py b/summarize_local_baselines.py index c2eecba..af6a270 100644 --- a/summarize_local_baselines.py +++ b/summarize_local_baselines.py @@ -26,11 +26,17 @@ def main(): if args.dataset is not None and metrics.get("dataset") != args.dataset: raise ValueError(f"Dataset mismatch in {path}: {metrics.get('dataset')}") if args.paired_deltas: - signature = (metrics.get("dataset"), metrics.get("sessions"), metrics.get("checkpoint_epoch"), + # Candidates may select different best epochs on the same dev set. + # Match the maximum training budget, not the selected epoch. + training_epochs = metrics.get( + "checkpoint_training_epochs", + metrics.get("checkpoint_epoch"), + ) + signature = (metrics.get("dataset"), metrics.get("sessions"), training_epochs, metrics.get("checkpoint_train_sessions"), metrics.get("checkpoint_dev_sessions"), metrics.get("loss"), metrics.get("layer_fusion")) if comparison_signature is not None and signature != comparison_signature: - raise ValueError(f"Paired runs have different splits/epochs/loss/fusion: {path}") + raise ValueError(f"Paired runs have different splits/budget/loss/fusion: {path}") comparison_signature = signature score = metrics[args.protocol] first = score["horizons"][0] @@ -54,7 +60,8 @@ def main(): if args.paired_deltas: print("\nPaired deltas (percentage points; same column order; matching seeds only)") for treatment, control in (("audio", "baseline"), ("visual", "baseline"), - ("av", "baseline"), ("av", "audio")): + ("av", "baseline"), ("av", "audio"), + ("track_gru", "pooled")): left, right = by_seed.get(treatment, {}), by_seed.get(control, {}) seeds = sorted(left.keys() & right.keys()) if not seeds: diff --git a/test_local_turn_features.py b/test_local_turn_features.py index 03e7103..2955caf 100644 --- a/test_local_turn_features.py +++ b/test_local_turn_features.py @@ -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] diff --git a/third_party/README.md b/third_party/README.md new file mode 100644 index 0000000..ced1e49 --- /dev/null +++ b/third_party/README.md @@ -0,0 +1,15 @@ +# Third-party reference code + +Reference repositories are inspected here but are not vendored into this +project. They are not runtime dependencies. + +The track-GRU design was checked against the official EgoSpeak repository at +revision `b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb`: + +```bash +git clone https://github.com/jun297/EgoSpeak.git third_party/EgoSpeak +git -C third_party/EgoSpeak checkout b5bee35e796d8f2d6aa7da2bf1b07836270f8bfb +``` + +`third_party/EgoSpeak/` is intentionally ignored. No EgoSpeak source file is +copied or imported by the project implementation. diff --git a/train_turn_head.py b/train_turn_head.py index b09f6a1..0ccc9ae 100644 --- a/train_turn_head.py +++ b/train_turn_head.py @@ -33,6 +33,8 @@ from easycom_head import ( TRANSITION_OBJECTIVE_NAMES, TurnPredictionHead, average_precision_metrics, + checkpoint_head_config, + checkpoint_loss_name, egospeak_multihorizon_loss, four_state_multihorizon_loss, head_output_names_for_loss, @@ -67,6 +69,7 @@ def parse_args() -> argparse.Namespace: default="pooled", help=( "pooled keeps the existing baseline; frame-set-attention is F1; " + "track-gru-residual adds a zero-gated causal track correction; " "audio-person-frame-set and gated-person-state-adapter are " "current-sidecar F5/F6 screening variants, not full F5/F6." ), @@ -77,6 +80,31 @@ def parse_args() -> argparse.Namespace: ) parser.add_argument("--person-attention-dim", type=int, default=256) parser.add_argument("--person-attention-heads", type=int, default=4) + parser.add_argument( + "--person-history-frames", + type=int, + default=16, + help="Maximum causal observations per current face track (16 at 2 FPS = 8s).", + ) + parser.add_argument( + "--person-track-max-gap-seconds", + type=float, + default=1.0, + help="Reset a provided track when consecutive observations are farther apart.", + ) + parser.add_argument( + "--initialize-pooled-checkpoint", + type=Path, + help=( + "Required by track-gru-residual: copy the exact pooled G+A+V base " + "before training the zero-gated track correction." + ), + ) + parser.add_argument( + "--train-pooled-base", + action="store_true", + help="Also update the initialized pooled base; default keeps it frozen.", + ) parser.add_argument( "--dataset", choices=DATASET_NAMES, @@ -184,6 +212,10 @@ 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.person_history_frames <= 0: + raise ValueError("--person-history-frames must be positive") + if args.person_track_max_gap_seconds <= 0: + raise ValueError("--person-track-max-gap-seconds must be positive") if args.local_fusion != "pooled" and ( args.local_feature_root is None or args.local_modalities != "av" ): @@ -191,6 +223,16 @@ def _validate_args(args: argparse.Namespace) -> None: "Person-set --local-fusion modes require --local-feature-root and " "--local-modalities av" ) + if args.local_fusion == "track-gru-residual": + if args.initialize_pooled_checkpoint is None: + raise ValueError( + "track-gru-residual requires --initialize-pooled-checkpoint" + ) + elif args.initialize_pooled_checkpoint is not None or args.train_pooled_base: + raise ValueError( + "--initialize-pooled-checkpoint/--train-pooled-base require " + "--local-fusion track-gru-residual" + ) if args.epochs <= 0: raise ValueError("--epochs must be positive") if args.batch_size <= 0: @@ -252,6 +294,7 @@ def _load_examples( local_contract: dict[str, Any] | None = None, compute_transition_targets: bool = True, local_fusion: str = "pooled", + person_track_max_gap_seconds: float = 1.0, ) -> tuple[ torch.Tensor, torch.Tensor, @@ -281,6 +324,7 @@ def _load_examples( local_modalities=local_modalities, local_contract=local_contract, local_fusion=local_fusion, + person_track_max_gap_seconds=person_track_max_gap_seconds, ) local_contract = data.feature_metadata.get("local_contract") states, targets, future_indices = data.future_examples(NUM_HORIZONS) @@ -467,7 +511,13 @@ def _evaluate_valid_horizons( person_mask = None if people is not None: rows = torch.arange(start, min(start + batch_size, len(states))) - person_batch, person_mask = people.padded(rows) + if head.local_fusion == "track-gru-residual": + person_batch, person_mask = people.padded_history( + rows, + max_observations=head.person_history_frames, + ) + else: + person_batch, person_mask = people.padded(rows) person_batch = person_batch.to(device=device, dtype=torch.float32) person_mask = person_mask.to(device=device) chunks.append( @@ -510,7 +560,7 @@ def _checkpoint( history: list[dict[str, Any]], ) -> dict[str, Any]: return { - "format_version": 6, + "format_version": 7, "head": head.state_dict(), "head_config": { "hidden_size": hidden_size, @@ -527,6 +577,8 @@ def _checkpoint( "local_fusion": head.local_fusion, "person_attention_dim": head.person_attention_dim, "person_attention_heads": head.person_attention_heads, + "person_history_frames": head.person_history_frames, + "person_track_max_gap_seconds": head.person_track_max_gap_seconds, }, "local_contract": feature_metadata[0].get("local_contract"), "alignment": { @@ -576,6 +628,14 @@ def _checkpoint( "local_projection_dim": args.local_projection_dim, "person_attention_dim": args.person_attention_dim, "person_attention_heads": args.person_attention_heads, + "person_history_frames": args.person_history_frames, + "person_track_max_gap_seconds": args.person_track_max_gap_seconds, + "initialize_pooled_checkpoint": ( + str(args.initialize_pooled_checkpoint.expanduser().resolve()) + if args.initialize_pooled_checkpoint is not None + else None + ), + "train_pooled_base": args.train_pooled_base, "learnable_layer_mixture": args.layer_fusion == "scalar-mixture", "supervision_policy": ( "transition-sparse-manifest" @@ -591,6 +651,94 @@ def _checkpoint( } +def _initialize_track_residual_from_pooled( + head: TurnPredictionHead, + checkpoint: Path, + *, + args: argparse.Namespace, + train_sessions: list[str], + dev_sessions: list[str], + local_contract: dict[str, Any] | None, +) -> None: + """Load and optionally freeze the exact pooled base inside a track model.""" + + path = checkpoint.expanduser().resolve() + if not path.is_file(): + raise FileNotFoundError(path) + payload = torch.load(path, map_location="cpu", weights_only=False) + if not isinstance(payload, dict) or not isinstance(payload.get("head"), dict): + raise ValueError(f"Invalid pooled checkpoint: {path}") + config = checkpoint_head_config(payload) + if config["local_fusion"] != "pooled": + raise ValueError( + "Track residual initialization requires a pooled checkpoint, got " + f"{config['local_fusion']!r}: {path}" + ) + if checkpoint_loss_name(payload) != args.loss: + raise ValueError("Pooled checkpoint loss differs from the track experiment") + if payload.get("dataset", args.dataset) != args.dataset: + raise ValueError("Pooled checkpoint dataset differs from the track experiment") + if payload.get("train_sessions") != train_sessions: + raise ValueError("Pooled checkpoint train sessions differ from this run") + if payload.get("dev_sessions", []) != dev_sessions: + raise ValueError("Pooled checkpoint dev sessions differ from this run") + if payload.get("local_contract") != local_contract: + raise ValueError("Pooled checkpoint local feature contract differs from this run") + + expected = { + "hidden_size": head.hidden_size, + "num_horizons": head.num_horizons, + "num_classes": head.num_classes, + "num_input_layers": head.num_input_layers, + "layer_fusion": head.layer_fusion, + "local_visual_dim": head.local_visual_dim, + "local_audio_dim": head.local_audio_dim, + "local_projection_dim": head.local_projection_dim, + } + mismatches = { + key: (config.get(key), value) + for key, value in expected.items() + if config.get(key) != value + } + if mismatches: + raise ValueError( + f"Pooled checkpoint architecture differs from track base: {mismatches}" + ) + + def is_base_parameter(name: str) -> bool: + return name == "layer_logits" or name.startswith( + ("layer_fusion_net.", "local_projections.", "net.") + ) + + source = payload["head"] + current = head.state_dict() + required = [name for name in current if is_base_parameter(name)] + missing = [name for name in required if name not in source] + bad_shapes = [ + name + for name in required + if name in source and source[name].shape != current[name].shape + ] + if missing or bad_shapes: + raise ValueError( + "Pooled checkpoint cannot initialize the exact base; " + f"missing={missing[:8]}, shape_mismatch={bad_shapes[:8]}" + ) + for name in required: + current[name] = source[name] + head.load_state_dict(current) + + if not args.train_pooled_base: + for name, parameter in head.named_parameters(): + if is_base_parameter(name): + parameter.requires_grad_(False) + print( + f"initialized pooled base: {path}; " + f"frozen={not args.train_pooled_base}", + flush=True, + ) + + def main() -> None: args = parse_args() _validate_args(args) @@ -679,6 +827,7 @@ def main() -> None: args.local_modalities, compute_transition_targets=args.transition_objective != "none", local_fusion=args.local_fusion, + person_track_max_gap_seconds=args.person_track_max_gap_seconds, ) local_contract = feature_metadata[0].get("local_contract") if not len(states): @@ -732,6 +881,7 @@ def main() -> None: local_contract, compute_transition_targets=args.transition_objective != "none", local_fusion=args.local_fusion, + person_track_max_gap_seconds=args.person_track_max_gap_seconds, ) if dev_hidden_size != hidden_size: raise ValueError( @@ -778,8 +928,20 @@ def main() -> None: local_fusion=args.local_fusion, person_attention_dim=args.person_attention_dim, person_attention_heads=args.person_attention_heads, + person_history_frames=args.person_history_frames, + person_track_max_gap_seconds=args.person_track_max_gap_seconds, **local_head_dimensions(local_contract), ).to(device) + if args.local_fusion == "track-gru-residual": + assert args.initialize_pooled_checkpoint is not None + _initialize_track_residual_from_pooled( + head, + args.initialize_pooled_checkpoint, + args=args, + train_sessions=train_sessions, + dev_sessions=dev_sessions, + local_contract=local_contract, + ) if local_contract is not None: print( f"local fusion: {args.local_fusion}/{args.local_modalities}; " @@ -861,6 +1023,13 @@ def main() -> None: best_epoch = args.epochs if not dev_sessions else 0 for epoch in range(1, args.epochs + 1): head.train() + if args.local_fusion == "track-gru-residual" and not args.train_pooled_base: + # Keep dropout/norm behavior in the frozen pooled checkpoint fixed + # while the new track branch learns a residual correction. + if head.layer_fusion_net is not None: + head.layer_fusion_net.eval() + head.local_projections.eval() + head.net.eval() loss_sum = 0.0 activity_loss_sum = 0.0 transition_loss_sum = 0.0 @@ -900,9 +1069,15 @@ def main() -> None: person_batch = None person_mask = None if person_features is not None: - person_batch, person_mask = person_features.padded( - example_index_batch - ) + if head.local_fusion == "track-gru-residual": + person_batch, person_mask = person_features.padded_history( + example_index_batch, + max_observations=head.person_history_frames, + ) + else: + person_batch, person_mask = person_features.padded( + example_index_batch + ) person_batch = person_batch.to( device=device, dtype=torch.float32,