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