Add causal person track GRU residual experiment
This commit is contained in:
@@ -11,6 +11,7 @@ __pycache__/
|
|||||||
/manifests/
|
/manifests/
|
||||||
/outputs/
|
/outputs/
|
||||||
/target_frame/
|
/target_frame/
|
||||||
|
/third_party/EgoSpeak/
|
||||||
|
|
||||||
# Local editors and environments.
|
# Local editors and environments.
|
||||||
.idea/
|
.idea/
|
||||||
|
|||||||
+22
-6
@@ -43,17 +43,27 @@ result.
|
|||||||
- Ego4D batch and 8-GPU local-feature launchers, with strict frame audits.
|
- Ego4D batch and 8-GPU local-feature launchers, with strict frame audits.
|
||||||
- Timing-derived wearer floor-state and event diagnostics.
|
- 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
|
### 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
|
- Two current-sidecar screening controls: `audio-person-frame-set` adds a
|
||||||
zero-gated audio-to-person read before the F1 query, and
|
zero-gated audio-to-person read before the F1 query, and
|
||||||
`gated-person-state-adapter` adds a zero-gated person residual after the
|
`gated-person-state-adapter` adds a zero-gated person residual after the
|
||||||
frozen Thinker readout. They are explicitly recorded as F5/F6 screening,
|
frozen Thinker readout. They are explicitly recorded as F5/F6 screening,
|
||||||
not full history-conditioned F5 or an internal Thinker F6 adapter.
|
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
|
### Current blocker
|
||||||
|
|
||||||
@@ -69,12 +79,17 @@ semantic event ground truth.
|
|||||||
|
|
||||||
### Next experiment gates
|
### 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.
|
1. **P0:** fix/freeze event attribution, splits and scenario definitions.
|
||||||
2. **P1:** probe aggregate activity and per-person active-speaker grounding;
|
2. **P1:** probe aggregate activity and per-person active-speaker grounding;
|
||||||
compare correct, wrong-person and time-shifted visual inputs.
|
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.
|
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,
|
4. **P3:** run the matched pooled/anonymous-track-GRU residual screen, then add
|
||||||
then add spatial or track-temporal structure only after its controls pass.
|
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.
|
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
|
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 |
|
| `infer_turn_head.py` | Save EgoSpeak-compatible predictions and AP/mAP |
|
||||||
| `evaluate_floor_state_changes.py` | CPU floor-state/event proxy evaluation |
|
| `evaluate_floor_state_changes.py` | CPU floor-state/event proxy evaluation |
|
||||||
| `summarize_local_baselines.py` | Aggregate local baseline `metrics.json` files |
|
| `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 |
|
| `train_omni_turn_lora.py` | Thinker LoRA training from cached AV |
|
||||||
| `eval_omni_turn_lora.py` | Thinker LoRA evaluation |
|
| `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
|
kept in a separate JSONL until pilot precision is accepted; they are not yet
|
||||||
silently merged into the training arrays.
|
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
|
## Candidate improvement: explicit working and long memory
|
||||||
|
|
||||||
Motivated by EgoSpeak's separation of long-term and working memory, test an
|
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 |
|
| [`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_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 |
|
| [`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
|
## 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 后
|
融合实验从 F1 Frame-set Attention 开始。后续模块只有在前一级通过预先定义的 gate 后
|
||||||
才实施。
|
才实施。
|
||||||
|
|
||||||
实现状态(2026-09-02):F1 的 loader、head、train/infer 路径和回归测试已经加入代码,
|
实现状态(更新于 2026-09-07):F1 已完成三 seed、四 epoch 的远程筛选,得到
|
||||||
但尚未在远程 PyTorch/L20 环境完成 runtime smoke,不能写成已验证结果。另加入两个仅用
|
`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` 和
|
当前 sidecar 的 F5/F6 **screening** 模式:`audio-person-frame-set` 和
|
||||||
`gated-person-state-adapter`。前者没有 F4 人物历史,后者位于 Thinker readout 之后;二者
|
`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. 不可改变的实验条件
|
## 2. 不可改变的实验条件
|
||||||
|
|
||||||
@@ -193,6 +196,14 @@ floor-action 结果。
|
|||||||
对照包括 frame-set、无 ID temporal pool、正确 track、打断连续性的 track 以及 past-shift。
|
对照包括 frame-set、无 ID temporal pool、正确 track、打断连续性的 track 以及 past-shift。
|
||||||
正确 track 必须优于无 ID 和打断 track,才能说明身份连续性有价值。
|
正确 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
|
## 8. F4:Person-conditioned audio/Omni history
|
||||||
|
|
||||||
F4 让每个匿名人物分别读取已有的 5 Hz 因果历史:
|
F4 让每个匿名人物分别读取已有的 5 Hz 因果历史:
|
||||||
@@ -268,11 +279,11 @@ README 必须写明 split、seed、checkpoint、输入特征版本、评测协
|
|||||||
|
|
||||||
## 12. 最近执行顺序
|
## 12. 最近执行顺序
|
||||||
|
|
||||||
1. 完成并测试 P0 competing-risk censoring,人工打印至少一条 response 和 yield 时间轴。
|
1. 用户当前选择先做 dense-AP model screen;运行 matched pooled 与
|
||||||
2. 扩展 local sidecar loader,使其在不破坏现有 pooled 加载路径的情况下返回逐帧人物
|
`track-gru-residual`,固定 Session 12 dev,不用测试集选历史长度。
|
||||||
feature、offset 和 mask;测试成功、失败、空人物和越界路径。
|
2. 若 track GRU 有稳定增益,补 broken-track、past-shift、no-ID 和 remove-person 控制;若
|
||||||
3. 实现 F1 最小 Frame-set Attention,不加入 bbox 或 track temporal。
|
无增益,只在 dev 上比较 8/16/32 observations 后停止该分支。
|
||||||
4. 单 session/短片 smoke:检查 shape、cutoff、chosen frame time、mask、空人物和排列
|
3. correct track 通过控制后,再实现 F4 person-conditioned audio/Omni history;否则回到
|
||||||
不变性。
|
对齐审计,不直接堆更深 interaction。
|
||||||
5. 在固定 train-internal dev 上运行 B1、B3、F1-a、F1-c、F1-d、F1-e 和 F1-g。
|
4. 在形成论文事件结论前,仍须完成 P0 competing-risk censoring,并人工打印 response 与
|
||||||
6. F1 通过后才进入 F2;若未通过,优先审计人物 crop 和 AV 对齐,不继续堆模型。
|
yield 时间轴。当前事件结果只能称为 timing-derived proxy。
|
||||||
|
|||||||
+191
-9
@@ -61,17 +61,23 @@ LAYER_FUSION_MODES = (
|
|||||||
"scalar-mixture",
|
"scalar-mixture",
|
||||||
"concat-mlp",
|
"concat-mlp",
|
||||||
)
|
)
|
||||||
PERSON_SET_FUSION_MODES = (
|
CURRENT_FRAME_PERSON_SET_FUSION_MODES = (
|
||||||
"frame-set-attention",
|
"frame-set-attention",
|
||||||
"audio-person-frame-set",
|
"audio-person-frame-set",
|
||||||
"gated-person-state-adapter",
|
"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_MODES = ("pooled", *PERSON_SET_FUSION_MODES)
|
||||||
LOCAL_FUSION_SCOPES = {
|
LOCAL_FUSION_SCOPES = {
|
||||||
"pooled": "pooled-baseline",
|
"pooled": "pooled-baseline",
|
||||||
"frame-set-attention": "f1-current-frame-person-set",
|
"frame-set-attention": "f1-current-frame-person-set",
|
||||||
"audio-person-frame-set": "f5-screening-current-frame-no-person-history",
|
"audio-person-frame-set": "f5-screening-current-frame-no-person-history",
|
||||||
"gated-person-state-adapter": "f6-screening-post-thinker-not-internal",
|
"gated-person-state-adapter": "f6-screening-post-thinker-not-internal",
|
||||||
|
"track-gru-residual": "causal-anonymous-track-gru-pooled-baseline-residual",
|
||||||
}
|
}
|
||||||
DatasetName = Literal["easycom", "ego4d"]
|
DatasetName = Literal["easycom", "ego4d"]
|
||||||
|
|
||||||
@@ -95,6 +101,8 @@ class TurnPredictionHead(nn.Module):
|
|||||||
local_fusion: str = "pooled",
|
local_fusion: str = "pooled",
|
||||||
person_attention_dim: int = 256,
|
person_attention_dim: int = 256,
|
||||||
person_attention_heads: int = 4,
|
person_attention_heads: int = 4,
|
||||||
|
person_history_frames: int = 16,
|
||||||
|
person_track_max_gap_seconds: float = 1.0,
|
||||||
) -> None:
|
) -> None:
|
||||||
super().__init__()
|
super().__init__()
|
||||||
if num_input_layers <= 0:
|
if num_input_layers <= 0:
|
||||||
@@ -143,6 +151,10 @@ class TurnPredictionHead(nn.Module):
|
|||||||
raise ValueError("Person attention dimension/heads must be positive")
|
raise ValueError("Person attention dimension/heads must be positive")
|
||||||
if person_attention_dim % person_attention_heads:
|
if person_attention_dim % person_attention_heads:
|
||||||
raise ValueError("Person attention dimension must be divisible by 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 (
|
if local_fusion in PERSON_SET_FUSION_MODES and (
|
||||||
local_visual_dim <= 0 or local_audio_dim <= 0
|
local_visual_dim <= 0 or local_audio_dim <= 0
|
||||||
):
|
):
|
||||||
@@ -152,6 +164,8 @@ class TurnPredictionHead(nn.Module):
|
|||||||
self.local_fusion = local_fusion
|
self.local_fusion = local_fusion
|
||||||
self.person_attention_dim = person_attention_dim
|
self.person_attention_dim = person_attention_dim
|
||||||
self.person_attention_heads = person_attention_heads
|
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:
|
if self.learnable_layer_mixture:
|
||||||
# Zero initialization exactly reproduces a uniform layer mean at
|
# Zero initialization exactly reproduces a uniform layer mean at
|
||||||
# step zero while allowing training to select useful query layers.
|
# step zero while allowing training to select useful query layers.
|
||||||
@@ -171,7 +185,7 @@ class TurnPredictionHead(nn.Module):
|
|||||||
fused_size = hidden_size
|
fused_size = hidden_size
|
||||||
projected_modalities = (
|
projected_modalities = (
|
||||||
(("audio", local_audio_dim),)
|
(("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))
|
else (("visual", local_visual_dim), ("audio", local_audio_dim))
|
||||||
)
|
)
|
||||||
for name, dimension in projected_modalities:
|
for name, dimension in projected_modalities:
|
||||||
@@ -195,6 +209,9 @@ class TurnPredictionHead(nn.Module):
|
|||||||
self.audio_person_gate: nn.Parameter | None = None
|
self.audio_person_gate: nn.Parameter | None = None
|
||||||
self.person_state_adapter: nn.Module | None = None
|
self.person_state_adapter: nn.Module | None = None
|
||||||
self.person_state_adapter_gate: nn.Parameter | 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:
|
if local_fusion in PERSON_SET_FUSION_MODES:
|
||||||
self.person_input_projection = nn.Sequential(
|
self.person_input_projection = nn.Sequential(
|
||||||
nn.LayerNorm(local_visual_dim),
|
nn.LayerNorm(local_visual_dim),
|
||||||
@@ -244,6 +261,25 @@ class TurnPredictionHead(nn.Module):
|
|||||||
nn.Linear(person_attention_dim, hidden_size),
|
nn.Linear(person_attention_dim, hidden_size),
|
||||||
)
|
)
|
||||||
self.person_state_adapter_gate = nn.Parameter(torch.zeros(()))
|
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:
|
else:
|
||||||
fused_size += person_attention_dim + 1
|
fused_size += person_attention_dim + 1
|
||||||
self.net = nn.Sequential(
|
self.net = nn.Sequential(
|
||||||
@@ -291,7 +327,7 @@ class TurnPredictionHead(nn.Module):
|
|||||||
|
|
||||||
@property
|
@property
|
||||||
def local_input_size(self) -> int:
|
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 self.local_audio_dim + 1
|
||||||
return sum(d + 1 for d in (self.local_visual_dim, self.local_audio_dim) if d)
|
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)
|
person_valid = person_mask.any(dim=1, keepdim=True).to(states.dtype)
|
||||||
return audio_projected, audio_valid, context[:, 0, :], person_valid
|
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(
|
def _shared_features(
|
||||||
self,
|
self,
|
||||||
states: torch.Tensor,
|
states: torch.Tensor,
|
||||||
@@ -379,7 +504,7 @@ class TurnPredictionHead(nn.Module):
|
|||||||
"""Return the activity MLP representation shared by both readouts."""
|
"""Return the activity MLP representation shared by both readouts."""
|
||||||
|
|
||||||
states = self._fuse_layers(states)
|
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:
|
if local is None:
|
||||||
raise ValueError("Person-set fusion requires local audio features")
|
raise ValueError("Person-set fusion requires local audio features")
|
||||||
audio, audio_valid, person_context, person_valid = self._frame_set_context(
|
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")
|
raise ValueError("Person sets supplied to a non-frame-set head")
|
||||||
return self.net[:-1](states)
|
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(
|
def forward(
|
||||||
self,
|
self,
|
||||||
states: torch.Tensor,
|
states: torch.Tensor,
|
||||||
@@ -435,13 +596,12 @@ class TurnPredictionHead(nn.Module):
|
|||||||
person_features: torch.Tensor | None = None,
|
person_features: torch.Tensor | None = None,
|
||||||
person_mask: torch.Tensor | None = None,
|
person_mask: torch.Tensor | None = None,
|
||||||
) -> torch.Tensor:
|
) -> torch.Tensor:
|
||||||
features = self._shared_features(
|
features, logits = self._activity_logits(
|
||||||
states,
|
states,
|
||||||
local,
|
local,
|
||||||
person_features,
|
person_features,
|
||||||
person_mask,
|
person_mask,
|
||||||
)
|
)
|
||||||
logits = self.net[-1](features)
|
|
||||||
return logits.view(
|
return logits.view(
|
||||||
*features.shape[:-1], self.num_horizons, self.num_classes
|
*features.shape[:-1], self.num_horizons, self.num_classes
|
||||||
)
|
)
|
||||||
@@ -459,13 +619,13 @@ class TurnPredictionHead(nn.Module):
|
|||||||
raise RuntimeError(
|
raise RuntimeError(
|
||||||
"This head was created without transition_auxiliary=True"
|
"This head was created without transition_auxiliary=True"
|
||||||
)
|
)
|
||||||
features = self._shared_features(
|
features, flat_activity_logits = self._activity_logits(
|
||||||
states,
|
states,
|
||||||
local,
|
local,
|
||||||
person_features,
|
person_features,
|
||||||
person_mask,
|
person_mask,
|
||||||
)
|
)
|
||||||
activity_logits = self.net[-1](features).view(
|
activity_logits = flat_activity_logits.view(
|
||||||
*features.shape[:-1],
|
*features.shape[:-1],
|
||||||
self.num_horizons,
|
self.num_horizons,
|
||||||
self.num_classes,
|
self.num_classes,
|
||||||
@@ -492,6 +652,10 @@ class TurnPredictionHead(nn.Module):
|
|||||||
result["person_state_adapter"] = float(
|
result["person_state_adapter"] = float(
|
||||||
torch.tanh(self.person_state_adapter_gate.detach()).cpu()
|
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
|
return result
|
||||||
|
|
||||||
|
|
||||||
@@ -1004,6 +1168,7 @@ def load_session_data(
|
|||||||
local_modalities: str = "av",
|
local_modalities: str = "av",
|
||||||
local_contract: dict[str, Any] | None = None,
|
local_contract: dict[str, Any] | None = None,
|
||||||
local_fusion: str = "pooled",
|
local_fusion: str = "pooled",
|
||||||
|
person_track_max_gap_seconds: float = 1.0,
|
||||||
) -> SessionData:
|
) -> SessionData:
|
||||||
"""Load and validate one saved Omni feature payload and target array."""
|
"""Load and validate one saved Omni feature payload and target array."""
|
||||||
|
|
||||||
@@ -1110,7 +1275,22 @@ def load_session_data(
|
|||||||
name,
|
name,
|
||||||
dataset=dataset,
|
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":
|
if local_modalities != "av":
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
"Person-set fusion requires --local-modalities av"
|
"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("local_fusion", "pooled")
|
||||||
config.setdefault("person_attention_dim", 256)
|
config.setdefault("person_attention_dim", 256)
|
||||||
config.setdefault("person_attention_heads", 4)
|
config.setdefault("person_attention_heads", 4)
|
||||||
|
config.setdefault("person_history_frames", 16)
|
||||||
|
config.setdefault("person_track_max_gap_seconds", 1.0)
|
||||||
config.setdefault(
|
config.setdefault(
|
||||||
"layer_fusion",
|
"layer_fusion",
|
||||||
(
|
(
|
||||||
|
|||||||
@@ -149,6 +149,12 @@ def _predict_states(
|
|||||||
start,
|
start,
|
||||||
min(start + batch_size, len(states)),
|
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_mask = local_person_features.padded(rows)
|
||||||
person_batch = person_batch.to(
|
person_batch = person_batch.to(
|
||||||
device=device,
|
device=device,
|
||||||
@@ -440,6 +446,7 @@ def main() -> None:
|
|||||||
local_modalities=local_contract["modalities"] if local_contract else "av",
|
local_modalities=local_contract["modalities"] if local_contract else "av",
|
||||||
local_contract=local_contract,
|
local_contract=local_contract,
|
||||||
local_fusion=head.local_fusion,
|
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"]):
|
if data.hidden_size != int(head_config["hidden_size"]):
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
@@ -519,12 +526,18 @@ def main() -> None:
|
|||||||
"prediction_layout": "[target_time, class, future_horizon]",
|
"prediction_layout": "[target_time, class, future_horizon]",
|
||||||
"checkpoint": str(args.checkpoint.resolve()),
|
"checkpoint": str(args.checkpoint.resolve()),
|
||||||
"checkpoint_epoch": checkpoint_payload.get("epoch"),
|
"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_train_sessions": checkpoint_payload.get("train_sessions"),
|
||||||
"checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
|
"checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
|
||||||
"layer_fusion": head.layer_fusion,
|
"layer_fusion": head.layer_fusion,
|
||||||
"local_fusion": head.local_fusion,
|
"local_fusion": head.local_fusion,
|
||||||
"local_fusion_scope": LOCAL_FUSION_SCOPES[head.local_fusion],
|
"local_fusion_scope": LOCAL_FUSION_SCOPES[head.local_fusion],
|
||||||
"person_interaction_gates": head.person_interaction_gates(),
|
"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,
|
"local_contract": local_contract,
|
||||||
"transition_auxiliary": head.transition_auxiliary,
|
"transition_auxiliary": head.transition_auxiliary,
|
||||||
"transition_objective": checkpoint_payload.get(
|
"transition_objective": checkpoint_payload.get(
|
||||||
|
|||||||
@@ -3,6 +3,8 @@
|
|||||||
- 会议:Findings of NAACL 2025
|
- 会议:Findings of NAACL 2025
|
||||||
- 阅读深度:深入阅读
|
- 阅读深度:深入阅读
|
||||||
- 官方链接:https://aclanthology.org/2025.findings-naacl.163/
|
- 官方链接: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 的收益很小,说明主要瓶颈可能不是输出空间。
|
它定义了本项目复现的核心协议: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
|
frame_offsets: torch.LongTensor
|
||||||
row_to_frame: torch.LongTensor
|
row_to_frame: torch.LongTensor
|
||||||
chosen_frame_times: torch.Tensor
|
chosen_frame_times: torch.Tensor
|
||||||
|
previous_observation: torch.LongTensor | None = None
|
||||||
|
|
||||||
def __post_init__(self) -> None:
|
def __post_init__(self) -> None:
|
||||||
if (
|
if (
|
||||||
@@ -65,6 +66,19 @@ class FramePersonFeatures:
|
|||||||
raise ValueError("Only rows without an available frame may have NaN frame time")
|
raise ValueError("Only rows without an available frame may have NaN frame time")
|
||||||
if not torch.isfinite(self.chosen_frame_times[~missing]).all():
|
if not torch.isfinite(self.chosen_frame_times[~missing]).all():
|
||||||
raise ValueError("Chosen person frame times must be finite")
|
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:
|
def __len__(self) -> int:
|
||||||
return len(self.row_to_frame)
|
return len(self.row_to_frame)
|
||||||
@@ -89,13 +103,14 @@ class FramePersonFeatures:
|
|||||||
self.frame_offsets,
|
self.frame_offsets,
|
||||||
self.row_to_frame[indices],
|
self.row_to_frame[indices],
|
||||||
self.chosen_frame_times[indices],
|
self.chosen_frame_times[indices],
|
||||||
|
self.previous_observation,
|
||||||
)
|
)
|
||||||
|
|
||||||
def padded(
|
def _padded_observation_indices(
|
||||||
self,
|
self,
|
||||||
row_indices: torch.Tensor | None = None,
|
row_indices: torch.Tensor | None = None,
|
||||||
) -> tuple[torch.Tensor, torch.BoolTensor]:
|
) -> tuple[torch.LongTensor, torch.BoolTensor]:
|
||||||
"""Materialize only one batch as [B,P,D] plus a True-for-real mask."""
|
"""Return current-frame observation indices and a real-person mask."""
|
||||||
|
|
||||||
if row_indices is None:
|
if row_indices is None:
|
||||||
frames = self.row_to_frame
|
frames = self.row_to_frame
|
||||||
@@ -121,12 +136,77 @@ class FramePersonFeatures:
|
|||||||
width = int(counts.max()) if len(counts) else 0
|
width = int(counts.max()) if len(counts) else 0
|
||||||
mask = torch.arange(width)[None, :] < counts[:, None]
|
mask = torch.arange(width)[None, :] < counts[:, None]
|
||||||
if width == 0:
|
if width == 0:
|
||||||
return self.features.new_zeros((len(frames), 0, self.dimension)), mask
|
return torch.zeros((len(frames), 0), dtype=torch.long), mask
|
||||||
person_indices = starts[:, None] + torch.arange(width)[None, :]
|
indices = starts[:, None] + torch.arange(width)[None, :]
|
||||||
safe_indices = person_indices.masked_fill(~mask, 0)
|
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]
|
values = self.features[safe_indices]
|
||||||
return values.masked_fill(~mask[..., None], 0), mask
|
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
|
@classmethod
|
||||||
def concatenate(
|
def concatenate(
|
||||||
cls,
|
cls,
|
||||||
@@ -143,6 +223,10 @@ class FramePersonFeatures:
|
|||||||
time_parts = []
|
time_parts = []
|
||||||
feature_base = 0
|
feature_base = 0
|
||||||
frame_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:
|
for part in parts:
|
||||||
feature_parts.append(part.features)
|
feature_parts.append(part.features)
|
||||||
offset_parts.append(part.frame_offsets[1:] + feature_base)
|
offset_parts.append(part.frame_offsets[1:] + feature_base)
|
||||||
@@ -151,6 +235,11 @@ class FramePersonFeatures:
|
|||||||
rows[valid] += frame_base
|
rows[valid] += frame_base
|
||||||
row_parts.append(rows)
|
row_parts.append(rows)
|
||||||
time_parts.append(part.chosen_frame_times)
|
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)
|
feature_base += len(part.features)
|
||||||
frame_base += len(part.frame_offsets) - 1
|
frame_base += len(part.frame_offsets) - 1
|
||||||
return cls(
|
return cls(
|
||||||
@@ -158,6 +247,7 @@ class FramePersonFeatures:
|
|||||||
torch.cat(offset_parts, dim=0),
|
torch.cat(offset_parts, dim=0),
|
||||||
torch.cat(row_parts, dim=0),
|
torch.cat(row_parts, dim=0),
|
||||||
torch.cat(time_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,
|
sample_id: str,
|
||||||
dataset: str,
|
dataset: str,
|
||||||
expected_contract: dict[str, Any] | None = None,
|
expected_contract: dict[str, Any] | None = None,
|
||||||
|
track_max_gap_seconds: float | None = None,
|
||||||
) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
|
) -> tuple[torch.Tensor, FramePersonFeatures, dict[str, Any]]:
|
||||||
"""Load direct audio and the latest causal per-frame person set per 5 Hz row."""
|
"""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]]
|
chosen[has_frame] = frame_times[row_to_frame[has_frame]]
|
||||||
if bool((chosen[has_frame] >= times[has_frame]).any()):
|
if bool((chosen[has_frame] >= times[has_frame]).any()):
|
||||||
raise ValueError(f"Person frame mapping selected future visual evidence: {path}")
|
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(
|
frame_sets = FramePersonFeatures(
|
||||||
features,
|
features,
|
||||||
frame_offsets,
|
frame_offsets,
|
||||||
row_to_frame.to(dtype=torch.long),
|
row_to_frame.to(dtype=torch.long),
|
||||||
chosen,
|
chosen,
|
||||||
|
previous_observation,
|
||||||
)
|
)
|
||||||
direct_audio = torch.cat(
|
direct_audio = torch.cat(
|
||||||
(audio.float(), audio_valid[:, None].float()),
|
(audio.float(), audio_valid[:, None].float()),
|
||||||
@@ -353,6 +456,96 @@ def load_frame_set_features(
|
|||||||
return direct_audio, frame_sets, contract
|
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]:
|
def local_head_dimensions(contract: dict[str, Any] | None) -> dict[str, int]:
|
||||||
if contract is None:
|
if contract is None:
|
||||||
return {"local_visual_dim": 0, "local_audio_dim": 0}
|
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:
|
if args.dataset is not None and metrics.get("dataset") != args.dataset:
|
||||||
raise ValueError(f"Dataset mismatch in {path}: {metrics.get('dataset')}")
|
raise ValueError(f"Dataset mismatch in {path}: {metrics.get('dataset')}")
|
||||||
if args.paired_deltas:
|
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("checkpoint_train_sessions"), metrics.get("checkpoint_dev_sessions"),
|
||||||
metrics.get("loss"), metrics.get("layer_fusion"))
|
metrics.get("loss"), metrics.get("layer_fusion"))
|
||||||
if comparison_signature is not None and signature != comparison_signature:
|
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
|
comparison_signature = signature
|
||||||
score = metrics[args.protocol]
|
score = metrics[args.protocol]
|
||||||
first = score["horizons"][0]
|
first = score["horizons"][0]
|
||||||
@@ -54,7 +60,8 @@ def main():
|
|||||||
if args.paired_deltas:
|
if args.paired_deltas:
|
||||||
print("\nPaired deltas (percentage points; same column order; matching seeds only)")
|
print("\nPaired deltas (percentage points; same column order; matching seeds only)")
|
||||||
for treatment, control in (("audio", "baseline"), ("visual", "baseline"),
|
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, {})
|
left, right = by_seed.get(treatment, {}), by_seed.get(control, {})
|
||||||
seeds = sorted(left.keys() & right.keys())
|
seeds = sorted(left.keys() & right.keys())
|
||||||
if not seeds:
|
if not seeds:
|
||||||
|
|||||||
+157
-2
@@ -33,7 +33,8 @@ from extract_turn_local_features import (
|
|||||||
)
|
)
|
||||||
from local_turn_features import (
|
from local_turn_features import (
|
||||||
LOCAL_FORMAT, FramePersonFeatures, causal_frame_ticks, crop_box, easycom_boxes,
|
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,
|
pooled_recent_visual, reference_alignment,
|
||||||
)
|
)
|
||||||
from train_turn_head import _load_examples
|
from train_turn_head import _load_examples
|
||||||
@@ -352,6 +353,72 @@ class LocalTests(unittest.TestCase):
|
|||||||
dataset="easycom",
|
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):
|
def test_frame_person_concatenation_selection_and_empty_padding(self):
|
||||||
features = torch.tensor([[1., 2.], [3., 4.]])
|
features = torch.tensor([[1., 2.], [3., 4.]])
|
||||||
a = FramePersonFeatures(
|
a = FramePersonFeatures(
|
||||||
@@ -567,6 +634,86 @@ class LocalTests(unittest.TestCase):
|
|||||||
restored.eval()
|
restored.eval()
|
||||||
self.assertTrue(torch.equal(restored(states, local, people, mask), logits))
|
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):
|
def test_head_shapes_gradients_missing_and_roundtrip(self):
|
||||||
states = torch.randn(3, 4, 8)
|
states = torch.randn(3, 4, 8)
|
||||||
for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)):
|
for v, a in ((0, 0), (6, 0), (0, 5), (6, 5)):
|
||||||
@@ -855,6 +1002,7 @@ class RoundTripTests(unittest.TestCase):
|
|||||||
"audio",
|
"audio",
|
||||||
"visual",
|
"visual",
|
||||||
"av",
|
"av",
|
||||||
|
"track-gru",
|
||||||
"frame-set",
|
"frame-set",
|
||||||
"f5-screen",
|
"f5-screen",
|
||||||
"f6-screen",
|
"f6-screen",
|
||||||
@@ -877,6 +1025,7 @@ class RoundTripTests(unittest.TestCase):
|
|||||||
if mode == "av":
|
if mode == "av":
|
||||||
train += ["--transition-objective", "r1-supervised"]
|
train += ["--transition-objective", "r1-supervised"]
|
||||||
person_mode = {
|
person_mode = {
|
||||||
|
"track-gru": "track-gru-residual",
|
||||||
"frame-set": "frame-set-attention",
|
"frame-set": "frame-set-attention",
|
||||||
"f5-screen": "audio-person-frame-set",
|
"f5-screen": "audio-person-frame-set",
|
||||||
"f6-screen": "gated-person-state-adapter",
|
"f6-screen": "gated-person-state-adapter",
|
||||||
@@ -887,10 +1036,16 @@ class RoundTripTests(unittest.TestCase):
|
|||||||
"--person-attention-dim", "4",
|
"--person-attention-dim", "4",
|
||||||
"--person-attention-heads", "2",
|
"--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":
|
if mode != "baseline":
|
||||||
modality = (
|
modality = (
|
||||||
"av"
|
"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
|
else mode
|
||||||
)
|
)
|
||||||
train += ["--local-feature-root", str(root / "local"), "--local-modalities", modality]
|
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,
|
TRANSITION_OBJECTIVE_NAMES,
|
||||||
TurnPredictionHead,
|
TurnPredictionHead,
|
||||||
average_precision_metrics,
|
average_precision_metrics,
|
||||||
|
checkpoint_head_config,
|
||||||
|
checkpoint_loss_name,
|
||||||
egospeak_multihorizon_loss,
|
egospeak_multihorizon_loss,
|
||||||
four_state_multihorizon_loss,
|
four_state_multihorizon_loss,
|
||||||
head_output_names_for_loss,
|
head_output_names_for_loss,
|
||||||
@@ -67,6 +69,7 @@ def parse_args() -> argparse.Namespace:
|
|||||||
default="pooled",
|
default="pooled",
|
||||||
help=(
|
help=(
|
||||||
"pooled keeps the existing baseline; frame-set-attention is F1; "
|
"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 "
|
"audio-person-frame-set and gated-person-state-adapter are "
|
||||||
"current-sidecar F5/F6 screening variants, not full F5/F6."
|
"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-dim", type=int, default=256)
|
||||||
parser.add_argument("--person-attention-heads", type=int, default=4)
|
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(
|
parser.add_argument(
|
||||||
"--dataset",
|
"--dataset",
|
||||||
choices=DATASET_NAMES,
|
choices=DATASET_NAMES,
|
||||||
@@ -184,6 +212,10 @@ def _validate_args(args: argparse.Namespace) -> None:
|
|||||||
raise ValueError("Person attention dimension/heads must be positive")
|
raise ValueError("Person attention dimension/heads must be positive")
|
||||||
if args.person_attention_dim % args.person_attention_heads:
|
if args.person_attention_dim % args.person_attention_heads:
|
||||||
raise ValueError("--person-attention-dim must be divisible by --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 (
|
if args.local_fusion != "pooled" and (
|
||||||
args.local_feature_root is None or args.local_modalities != "av"
|
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 "
|
"Person-set --local-fusion modes require --local-feature-root and "
|
||||||
"--local-modalities av"
|
"--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:
|
if args.epochs <= 0:
|
||||||
raise ValueError("--epochs must be positive")
|
raise ValueError("--epochs must be positive")
|
||||||
if args.batch_size <= 0:
|
if args.batch_size <= 0:
|
||||||
@@ -252,6 +294,7 @@ def _load_examples(
|
|||||||
local_contract: dict[str, Any] | None = None,
|
local_contract: dict[str, Any] | None = None,
|
||||||
compute_transition_targets: bool = True,
|
compute_transition_targets: bool = True,
|
||||||
local_fusion: str = "pooled",
|
local_fusion: str = "pooled",
|
||||||
|
person_track_max_gap_seconds: float = 1.0,
|
||||||
) -> tuple[
|
) -> tuple[
|
||||||
torch.Tensor,
|
torch.Tensor,
|
||||||
torch.Tensor,
|
torch.Tensor,
|
||||||
@@ -281,6 +324,7 @@ def _load_examples(
|
|||||||
local_modalities=local_modalities,
|
local_modalities=local_modalities,
|
||||||
local_contract=local_contract,
|
local_contract=local_contract,
|
||||||
local_fusion=local_fusion,
|
local_fusion=local_fusion,
|
||||||
|
person_track_max_gap_seconds=person_track_max_gap_seconds,
|
||||||
)
|
)
|
||||||
local_contract = data.feature_metadata.get("local_contract")
|
local_contract = data.feature_metadata.get("local_contract")
|
||||||
states, targets, future_indices = data.future_examples(NUM_HORIZONS)
|
states, targets, future_indices = data.future_examples(NUM_HORIZONS)
|
||||||
@@ -467,6 +511,12 @@ def _evaluate_valid_horizons(
|
|||||||
person_mask = None
|
person_mask = None
|
||||||
if people is not None:
|
if people is not None:
|
||||||
rows = torch.arange(start, min(start + batch_size, len(states)))
|
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_mask = people.padded(rows)
|
||||||
person_batch = person_batch.to(device=device, dtype=torch.float32)
|
person_batch = person_batch.to(device=device, dtype=torch.float32)
|
||||||
person_mask = person_mask.to(device=device)
|
person_mask = person_mask.to(device=device)
|
||||||
@@ -510,7 +560,7 @@ def _checkpoint(
|
|||||||
history: list[dict[str, Any]],
|
history: list[dict[str, Any]],
|
||||||
) -> dict[str, Any]:
|
) -> dict[str, Any]:
|
||||||
return {
|
return {
|
||||||
"format_version": 6,
|
"format_version": 7,
|
||||||
"head": head.state_dict(),
|
"head": head.state_dict(),
|
||||||
"head_config": {
|
"head_config": {
|
||||||
"hidden_size": hidden_size,
|
"hidden_size": hidden_size,
|
||||||
@@ -527,6 +577,8 @@ def _checkpoint(
|
|||||||
"local_fusion": head.local_fusion,
|
"local_fusion": head.local_fusion,
|
||||||
"person_attention_dim": head.person_attention_dim,
|
"person_attention_dim": head.person_attention_dim,
|
||||||
"person_attention_heads": head.person_attention_heads,
|
"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"),
|
"local_contract": feature_metadata[0].get("local_contract"),
|
||||||
"alignment": {
|
"alignment": {
|
||||||
@@ -576,6 +628,14 @@ def _checkpoint(
|
|||||||
"local_projection_dim": args.local_projection_dim,
|
"local_projection_dim": args.local_projection_dim,
|
||||||
"person_attention_dim": args.person_attention_dim,
|
"person_attention_dim": args.person_attention_dim,
|
||||||
"person_attention_heads": args.person_attention_heads,
|
"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",
|
"learnable_layer_mixture": args.layer_fusion == "scalar-mixture",
|
||||||
"supervision_policy": (
|
"supervision_policy": (
|
||||||
"transition-sparse-manifest"
|
"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:
|
def main() -> None:
|
||||||
args = parse_args()
|
args = parse_args()
|
||||||
_validate_args(args)
|
_validate_args(args)
|
||||||
@@ -679,6 +827,7 @@ def main() -> None:
|
|||||||
args.local_modalities,
|
args.local_modalities,
|
||||||
compute_transition_targets=args.transition_objective != "none",
|
compute_transition_targets=args.transition_objective != "none",
|
||||||
local_fusion=args.local_fusion,
|
local_fusion=args.local_fusion,
|
||||||
|
person_track_max_gap_seconds=args.person_track_max_gap_seconds,
|
||||||
)
|
)
|
||||||
local_contract = feature_metadata[0].get("local_contract")
|
local_contract = feature_metadata[0].get("local_contract")
|
||||||
if not len(states):
|
if not len(states):
|
||||||
@@ -732,6 +881,7 @@ def main() -> None:
|
|||||||
local_contract,
|
local_contract,
|
||||||
compute_transition_targets=args.transition_objective != "none",
|
compute_transition_targets=args.transition_objective != "none",
|
||||||
local_fusion=args.local_fusion,
|
local_fusion=args.local_fusion,
|
||||||
|
person_track_max_gap_seconds=args.person_track_max_gap_seconds,
|
||||||
)
|
)
|
||||||
if dev_hidden_size != hidden_size:
|
if dev_hidden_size != hidden_size:
|
||||||
raise ValueError(
|
raise ValueError(
|
||||||
@@ -778,8 +928,20 @@ def main() -> None:
|
|||||||
local_fusion=args.local_fusion,
|
local_fusion=args.local_fusion,
|
||||||
person_attention_dim=args.person_attention_dim,
|
person_attention_dim=args.person_attention_dim,
|
||||||
person_attention_heads=args.person_attention_heads,
|
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),
|
**local_head_dimensions(local_contract),
|
||||||
).to(device)
|
).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:
|
if local_contract is not None:
|
||||||
print(
|
print(
|
||||||
f"local fusion: {args.local_fusion}/{args.local_modalities}; "
|
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
|
best_epoch = args.epochs if not dev_sessions else 0
|
||||||
for epoch in range(1, args.epochs + 1):
|
for epoch in range(1, args.epochs + 1):
|
||||||
head.train()
|
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
|
loss_sum = 0.0
|
||||||
activity_loss_sum = 0.0
|
activity_loss_sum = 0.0
|
||||||
transition_loss_sum = 0.0
|
transition_loss_sum = 0.0
|
||||||
@@ -900,6 +1069,12 @@ def main() -> None:
|
|||||||
person_batch = None
|
person_batch = None
|
||||||
person_mask = None
|
person_mask = None
|
||||||
if person_features is not 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(
|
person_batch, person_mask = person_features.padded(
|
||||||
example_index_batch
|
example_index_batch
|
||||||
)
|
)
|
||||||
|
|||||||
Reference in New Issue
Block a user