Add reproducible EasyCom causal head pipeline

This commit is contained in:
Xin Wang
2026-08-16 22:19:43 +08:00
parent bde72487bb
commit 080c45f049
8 changed files with 1451 additions and 470 deletions
+81 -84
View File
@@ -32,9 +32,10 @@ classification head.
| `smoke_test.py` | Verified on one NVIDIA L20 | Text and 12-second EasyCom audio-video inference |
| `easycom_manifest.py` | Implemented, awaiting dataset-side run | Discover one-minute clips, preserve session timecodes, and check 5 Hz targets |
| `qwen_omni_features.py` | Implemented, not executed here | Load only the 3B Thinker and extract fused states from one AV window |
| `streaming_turn_model.py` | Prototype implemented, not executed here | Produce one causal Omni state every 200 ms, with reference and KV-cache modes |
| `streaming_turn_model.py` | Implemented | Produce one causal Omni state every 200 ms, with causal-prefix and open-stream modes |
| `TurnPredictionHead` | Defined, untrained | Map `[2048]` to `[10, 3]` logits |
| Head training and EgoSpeak evaluation script | Not yet implemented | Align saved states with `target_perframe`, train the head, and reproduce Table 2 mAP |
| `train_turn_head.py` | Implemented | Train the frozen-Omni anticipation head with the official EasyCom split |
| `infer_turn_head.py` | Implemented | Save EgoSpeak-format predictions and reproduce per-horizon AP/mAP |
The official full model, Thinker input path, embedded audio, video decoding,
and FlashAttention-2 have passed a 12-second EasyCom smoke test on one NVIDIA
@@ -284,13 +285,17 @@ Session_1.json clip order, exact T, crop and sampling provenance
Session_1.ffconcat reproducible dense concatenation order
```
The script fails rather than silently padding or truncating when
`sum(source_frames) / 4 != target.shape[0]`.
The target length must equal the number of complete 5 Hz intervals,
`floor(sum(source_frames) / 4)`. If a session ends with one to three residual
20 Hz frames, the target-derived duration discards only that incomplete final
interval; the discarded frame count is recorded in the sidecar.
Each clip's decoded video and audio are trimmed to the exact duration implied
by its video frame count before concatenation. This prevents AAC priming/tail
padding from accumulating at the 24 clip boundaries and causing progressive
audio/video drift.
audio/video drift. The session-global visual grid is padded defensively at EOF
and then trimmed by frame index, so it contains exactly
`ceil(target_steps / 5 * video_fps)` frames on every FFmpeg version.
On the remote L20 machine, use its own dataset paths. For example:
@@ -336,7 +341,7 @@ That early state can contain audio from 0.2--2.0 seconds. At prediction time
## Causal 5 Hz extraction modes
`streaming_turn_model.py` implements three modes.
`streaming_turn_model.py` implements two modes.
### 1. `causal-prefix`: correctness and semantic baseline
@@ -368,57 +373,7 @@ python streaming_turn_model.py \
--output features/easycom_20s_prefix.pt
```
### 2. `commit-replay`: closed-turn compatibility baseline
The extractor retains Thinker KV for completed two-second media blocks. During
the unfinished block it branches from the same committed cache and evaluates
only the available prefix:
```text
committed historical KV + current block [0.0, 0.2] -> q_t
committed historical KV + current block [0.0, 0.4] -> q_t
...
committed historical KV + current block [0.0, 2.0] -> q_t
-> commit media KV
```
The temporary readout/assistant query is never committed. Only the completed
media block enters long-term cache. This avoids recomputing the full history,
but still re-encodes at most the current two-second prefix.
Example with numerical checking:
```bash
python streaming_turn_model.py \
--video /path/to/easycom_session.mp4 \
--use-audio-in-video \
--mode commit-replay \
--start 0 \
--end 20 \
--memory-seconds 20 \
--verify-every 5 \
--verify-continuous-every 5 \
--output features/easycom_20s_stream.pt
```
`--verify-every 5` compares cached processing with a from-scratch replay of the
same sequence of two-second media turns. The saved cosine similarity should be
extremely close to 1. If it is not, the TMRoPE/cache path is invalid and must
not be used for a paper result.
`--verify-continuous-every 5` is a different check: it compares the cached
sequence of two-second media turns with a single continuous causal window.
This comparison need not be numerically identical, but it measures whether the
input organization materially changes the representation. Save and report the
two check types separately (`cache_replay` and `continuous_window_semantics`).
There is an important experimental distinction: `commit-replay` represents
history as a sequence of two-second media turns. Its numerical check proves
cache equivalence to that same turn sequence, not equivalence to one continuous
20-second video placeholder. Compare it against `causal-prefix` on identical
clips and report the representation if it changes mAP.
### 3. `open-stream-replay`: recommended efficient mode
### 2. `open-stream-replay`: recommended efficient mode
This mode keeps one user/media turn open across the whole bounded-memory
segment. The first block contributes the chat and AV opening markers once;
@@ -468,20 +423,6 @@ Twenty seconds is a practical memory/context setting for the 3090, not the
maximum video duration. The complete session is processed sequentially and a
prediction is emitted every 0.2 seconds.
The legacy `commit-replay` prototype resets Thinker KV at each non-overlapping
20-second boundary:
```text
session: 0 -------- 20 -------- 40 -------- 60 ...
cache: [segment 1] [segment 2] [segment 3] ...
output: continuous predictions across the entire session
```
This preserves complete output coverage but limits available history after a
reset. Report the memory length and additionally evaluate the first two seconds
after each reset versus all other timesteps. Rolling KV eviction is a possible
future optimization, but requires correct multimodal position rebasing.
`open-stream-replay` instead uses blockwise rolling memory. With a 20-second
budget it rebuilds the open cache from the latest nine completed blocks
(18 seconds) when a new block starts, then adds the current 0--2 second tail.
@@ -515,6 +456,10 @@ uses only frames with timestamps `<= t`, and duplicates the most recent past
frame when the visual tower needs an even frame count. Audio is loaded once at
16 kHz and sliced exactly at each causal boundary. The processor still receives
`use_audio_in_video=True`, so Qwen's audio/video token interleaving is retained.
Long runs use `tqdm` and update the progress bar after every completed state,
showing percentage, media time, elapsed time, ETA, and throughput. Pass
`--no-progress` to disable it. The Thinker-only path also suppresses Qwen's
irrelevant speech-output system prompt warning.
Before a full session, run only the first four seconds and inspect both check
types:
@@ -536,7 +481,7 @@ python streaming_turn_model.py \
Do not launch all sessions unless `cache_replay` cosine is nearly 1. The
`continuous_window_semantics` cosine answers a different question and may be
lower because media-turn boundaries change the representation.
lower because blockwise media preprocessing changes the representation.
Video can be sampled at 2 FPS while predictions remain at 5 Hz. Audio and the
causal boundary update every 200 ms; consecutive prediction times may use the
@@ -552,7 +497,7 @@ Each `.pt` file contains:
- `label_indices`: `[0, ..., N-1]`, direct `target_perframe` row alignment;
- `valid_10_horizon_steps`: `N-10`, because the last ten current rows do not
have all future labels;
- `mode`: `causal-prefix`, `commit-replay`, or `open-stream-replay`;
- `mode`: `causal-prefix` or `open-stream-replay`;
- `stream_representation`: explicit input-organization metadata;
- `config`: step, native block, memory, FPS, and pixel settings;
- `verification`: cached-versus-exact numerical checks;
@@ -575,7 +520,8 @@ LayerNorm(2048)
-> reshape [10, 3]
```
`streaming_turn_model.py` exposes:
`easycom_head.py` defines the head and loss; they remain re-exported by
`streaming_turn_model.py` for compatibility:
```python
from streaming_turn_model import (
@@ -584,16 +530,37 @@ from streaming_turn_model import (
)
```
Recommended initial procedure:
The supplied LSTR config uses Adam, 50 epochs, learning rate `7e-5`, weight
decay `5e-5`, dropout `0.2`, and 10 warmup epochs followed by cosine decay.
`train_turn_head.py` uses those defaults. Its batch size is 512 causal states,
not LSTR's 16 temporal windows, because the two model inputs have different
shapes.
Train once on the official nine sessions with a fixed 50-epoch protocol:
```bash
python train_turn_head.py \
--feature-root features/causal_prefix_2fps \
--target-root target_frame \
--output checkpoints/causal_prefix_2fps_head.pth \
--device cuda:0
```
The default loss is `lstr-mce`, exactly matching LSTR's MCE treatment of a
multi-positive `[0,1,1]` row: it sums both active negative-log-softmax terms
without normalizing the target. The previous L2-normalized objective belongs
to MiniROAD and is available only as the explicit
`--loss miniroad-normalized` ablation.
Recommended experimental procedure:
1. Freeze Qwen and extract all training-session states offline.
2. Verify state/label time offset and truncate only the final ten invalid
current positions.
3. Keep EgoSpeak targets as float multi-positive tensors `[N, 10, 3]`; never
convert them with `argmax`.
4. Train with `egospeak_multihorizon_loss`. It exactly mirrors MiniROAD's
anticipation objective: L2-normalize each three-class target, then compute
its weighted negative log-softmax.
4. For comparison with the supplied LSTR shell, train with the default
`egospeak_multihorizon_loss`/`lstr-mce` objective.
5. Select hyperparameters using only training sessions. For EasyCom, reserve a
training session for development or use session-level cross-validation;
never tune on Sessions 1--3.
@@ -605,7 +572,22 @@ ablations. They should not replace the frozen-Omni linear/MLP-head baseline.
## EgoSpeak-compatible evaluation
For every valid current timestep, produce probabilities shaped:
Run inference on the official test sessions (`Session_3`, `Session_2`, and
`Session_1`):
```bash
python infer_turn_head.py \
--feature-root features/causal_prefix_2fps \
--target-root target_frame \
--checkpoint checkpoints/causal_prefix_2fps_head.pth \
--output-dir results/causal_prefix_2fps_head \
--device cuda:0
```
This writes `pred/Session_X.npy` with the exact LSTR/EgoSpeak layout
`[target_time, class, future_horizon]`, i.e. `[T,3,10]`, plus `metrics.json`.
For every valid current timestep, the head itself produces probabilities
shaped:
```text
[number of timesteps, 10 future horizons, 3 classes]
@@ -621,8 +603,22 @@ Then reproduce EgoSpeak's Table 2 procedure:
5. average the ten horizon mAP values for the final score.
With EgoSpeak's `ignore_index=-100`, the released evaluator does not remove
class index zero, so background is included. Reuse the same implementation or
verify a reimplementation numerically.
class index zero, so background is included. The script reports two protocols:
- `egospeak_full_timeline`: the same full `[T,3,10]` array and evaluator as
LSTR, with zeros where a prediction would require pre-session context;
- `causal_valid_only`: excludes only those unavailable first `h` rows for
horizon `h`. This is the primary strictly causal score.
The difference involves at most the first ten rows of each session, but both
are saved so the comparison is explicit.
The convenience wrapper mirrors EgoSpeak's train/inference shells:
```bash
chmod +x run_easycom_head_experiment.sh
CUDA_VISIBLE_DEVICES=0 ./run_easycom_head_experiment.sh all
```
Reference Table 2 targets from the EgoSpeak paper:
@@ -701,7 +697,7 @@ pip install -r requirements.txt
```
The implementation is pinned to `transformers==4.52.3` because the
commit/replay path uses version-specific Thinker internals. On Linux,
open-stream replay path uses version-specific Thinker internals. On Linux,
`qwen-omni-utils[decord]` is recommended for faster video decoding. On a
compatible NVIDIA GPU, optionally install `flash-attn` and pass
`--flash-attention-2`.
@@ -732,8 +728,9 @@ Do not launch full EasyCom/Ego4D extraction until all of these pass:
time `t`;
- `states` and `target_perframe` lengths/time origins are reconciled and the
chosen label offset is saved;
- commit/replay cosine equivalence is near 1 on short clips;
- commit/replay and causal-prefix predictions are compared on identical clips;
- open-stream cache replay is exactly equivalent to its from-scratch replay on
short clips;
- open-stream and causal-prefix predictions are compared on identical clips;
- a 20-second run fits the 3090 with the selected pixel budget and attention
implementation;
- saved outputs contain exactly ten future horizons in the order 0.2--2.0 s;