Merge branch 'main' of https://gitea.xiaowang.eu.org/wx44wx/icassp2027-experiment
This commit is contained in:
Xin Wang
2
AGENTS.md
Normal file
2
AGENTS.md
Normal file
@@ -0,0 +1,2 @@
|
||||
作为一名资深音视频，深度学习国际知名导师。
|
||||
对于我的想法，引导我以思辨的方式继续，指导我完成icassp2027论文实验和编写。
|
||||
375
README.md
375
README.md
@@ -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:
|
||||
|
||||
@@ -312,33 +317,112 @@ script and reuse the existing MP4 without encoding it again:
|
||||
python prepare_easycom_session.py ... --fps 2 --reuse-existing
|
||||
```
|
||||
|
||||
## What Qwen's native 2-second chunk means
|
||||
## What Qwen's native 2-second streaming design means
|
||||
|
||||
Qwen2.5-Omni is **not restricted to receiving audio once every two seconds**.
|
||||
It can receive an arbitrary-duration AV input. The two-second value describes
|
||||
internal processing:
|
||||
Qwen2.5-Omni was designed for **block-wise multimodal prefill**. Its
|
||||
[technical report](https://arxiv.org/html/2503.20215) states that the audio
|
||||
and visual encoders are modified for block-wise processing and that audio
|
||||
self-attention is performed in independent two-second blocks. The released
|
||||
[3B configuration](https://huggingface.co/Qwen/Qwen2.5-Omni-3B/raw/main/config.json)
|
||||
uses `seconds_per_chunk=2`, `position_id_per_seconds=25`, and audio
|
||||
`n_window=100`.
|
||||
|
||||
1. the audio encoder divides acoustic features into approximately two-second
|
||||
attention blocks;
|
||||
2. audio and video tokens are interleaved in two-second temporal groups before
|
||||
entering the Thinker.
|
||||
|
||||
The audio encoder attention inside a block is non-causal. Therefore this is
|
||||
invalid for online evaluation:
|
||||
Conceptually, a native online input is organized as follows:
|
||||
|
||||
```text
|
||||
encode the complete [0, 2.0 s] block
|
||||
-> read a purported state at 0.2 s
|
||||
open one user/media turn
|
||||
encode/append AV [0,2] -> commit Thinker KV
|
||||
encode/append AV [2,4] -> commit Thinker KV
|
||||
encode/append AV [4,6] -> commit Thinker KV
|
||||
...
|
||||
close the media turn -> generate text and/or speech
|
||||
```
|
||||
|
||||
That early state can contain audio from 0.2--2.0 seconds. At prediction time
|
||||
`t`, the model must receive only audio and video with timestamps `<= t`.
|
||||
Within each temporal group, video tokens precede the aligned audio tokens.
|
||||
TMRoPE assigns 25 temporal positions per second, or 50 positions per native
|
||||
two-second group, so audio and video share a physical-time coordinate system.
|
||||
Completed blocks are appended to one causal Thinker sequence; the long-range
|
||||
history is modeled by the Thinker rather than by cross-block full attention in
|
||||
the modality encoders.
|
||||
|
||||
This architecture is block-causal, not frame-causal. Attention inside an
|
||||
audio block is non-causal. It is therefore invalid to encode the complete
|
||||
`[0,2]` block and claim that its representation at 0.2 seconds saw no future:
|
||||
|
||||
```text
|
||||
encode complete [0,2] AV -> read state at 0.2 s # future leakage
|
||||
```
|
||||
|
||||
For this project's 5 Hz targets, the current incomplete block must instead be
|
||||
encoded only through the prediction time. These 0.2-second partial-block
|
||||
readouts are an extension made for EgoSpeak; the Qwen paper does not define a
|
||||
canonical hidden state every 200 ms.
|
||||
|
||||
The public Transformers and local Qwen web-demo APIs normally receive a
|
||||
completed media object and call the processor/model once. They retain the
|
||||
model's block-wise masks and token organization, but do not expose the full
|
||||
persistent live-session implementation used by the hosted real-time product.
|
||||
|
||||
## Native Qwen streaming versus `causal-prefix`
|
||||
|
||||
The two extraction modes are both causal with respect to raw observations,
|
||||
but they define different model inputs. `causal-prefix` is a sliding-window
|
||||
reference invented for this experiment; it is not Qwen's native streaming
|
||||
execution model.
|
||||
|
||||
| Property | `causal-prefix` | Native Qwen streaming idea |
|
||||
|---|---|---|
|
||||
| Session | New inference at every 0.2 s | One open media turn |
|
||||
| Media organization | One complete `[max(0,t-M),t]` item | Fixed source-aligned 2 s blocks |
|
||||
| Modality encoding | Re-encode the complete window | Encode/commit each completed block once |
|
||||
| Thinker KV | Start empty for every state | Append to persistent KV |
|
||||
| TMRoPE | Recomputed from each window origin | Advances with the stream |
|
||||
| History | Exact rolling `M` seconds | Accumulates until a serving/context policy intervenes |
|
||||
| 5 Hz readout | Close/read a new turn each time | Not specified by the original model |
|
||||
|
||||
The difference becomes especially important after a sliding window fills. At
|
||||
`t=20.8` with `M=20`, `causal-prefix` constructs one new media item covering:
|
||||
|
||||
```text
|
||||
[0.8,20.8]
|
||||
-> encoder blocks [0.8,2.8], [2.8,4.8], ..., [18.8,20.8]
|
||||
```
|
||||
|
||||
At `t=21.0`, every encoder boundary moves again:
|
||||
|
||||
```text
|
||||
[1.0,21.0]
|
||||
-> encoder blocks [1.0,3.0], [3.0,5.0], ..., [19.0,21.0]
|
||||
```
|
||||
|
||||
Native streaming keeps the source-aligned partition instead:
|
||||
|
||||
```text
|
||||
[0,2], [2,4], ..., [18,20], current partial [20,20.8]
|
||||
```
|
||||
|
||||
Because Qwen's audio encoder uses block-local non-causal attention, moving a
|
||||
block boundary changes the representation even when most raw samples are the
|
||||
same. Window-relative TMRoPE, convolution/padding boundaries, chat closure,
|
||||
and BF16/attention execution add further differences. This is why exact KV
|
||||
replay can be numerically identical to its own from-scratch replay while still
|
||||
differing from `causal-prefix`; the modes are evaluating different token and
|
||||
encoder histories.
|
||||
|
||||
Use the modes according to the scientific question:
|
||||
|
||||
- use `open-stream-replay` for the primary online experiment, training, and
|
||||
deployment-aligned evaluation;
|
||||
- use `causal-prefix` as an expensive continuous-window reference and
|
||||
short-clip ablation;
|
||||
- never train a head on one representation and deploy it on the other without
|
||||
retraining, because their feature distributions are not interchangeable.
|
||||
|
||||
## 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
|
||||
### 1. `causal-prefix`: continuous-window reference
|
||||
|
||||
At every 200 ms, independently process the causal context ending at the
|
||||
current time:
|
||||
@@ -353,7 +437,8 @@ Properties:
|
||||
- one continuous media window per prediction;
|
||||
- no cross-forward cache;
|
||||
- very expensive because overlapping context is recomputed;
|
||||
- preferred reference when checking semantics and causality.
|
||||
- useful for checking continuous sliding-window semantics and causality;
|
||||
- not the native Qwen block-streaming execution model.
|
||||
|
||||
Example:
|
||||
|
||||
@@ -368,57 +453,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;
|
||||
@@ -435,6 +470,18 @@ open AV prefix
|
||||
AV/chat closure -> never committed; readout branch only
|
||||
```
|
||||
|
||||
This matches the core Qwen design choices: fixed two-second AV groups, one
|
||||
open media turn, completed-block commits, and incremental Thinker KV. It makes
|
||||
two practical changes for this dataset:
|
||||
|
||||
1. the current 0--2 second partial block is evaluated speculatively every
|
||||
0.2 seconds and discarded unless it reaches a complete block boundary;
|
||||
2. the KV history is bounded and rebuilt from recent complete blocks rather
|
||||
than accumulating for an entire multi-minute session.
|
||||
|
||||
It should therefore be described as **bounded Qwen-style open streaming**, not
|
||||
as bit-for-bit reproduction of an unbounded official serving session.
|
||||
|
||||
The audio encoder accepts a shorter final chunk with padding and an
|
||||
actual-length mask. For the two-frame visual patch requirement, strict-grid
|
||||
mode duplicates only the most recent past frame when necessary.
|
||||
@@ -465,29 +512,26 @@ required even though artificial user-turn boundaries have been removed.
|
||||
## Whole-session prediction and the 20-second limit
|
||||
|
||||
Twenty seconds is a practical memory/context setting for the 3090, not the
|
||||
maximum video duration. The complete session is processed sequentially and a
|
||||
maximum video duration and not a limit specified by Qwen. The original
|
||||
append-only idea eventually reaches both linear KV-memory growth and the
|
||||
Thinker's 32,768-token context limit, so an unbounded multi-minute stream is
|
||||
not practical. 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
|
||||
`open-stream-replay` 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.
|
||||
Its effective history therefore varies from 18 to 20 seconds and never drops
|
||||
to an empty 0.2-second context at a 20-second boundary.
|
||||
|
||||
For a short native-stream fidelity check, set `--memory-seconds` to at least
|
||||
the tested clip duration. For example, the options
|
||||
`--start 0 --end 40 --memory-seconds 40` prevent a rolling rebuild during that
|
||||
40-second clip.
|
||||
Do not extend this policy to complete EasyCom/Ego4D sessions: KV memory grows
|
||||
linearly, and this implementation clones the cache for speculative partial
|
||||
blocks and readout branches.
|
||||
|
||||
For the paper experiment, drive the output length from the preparation sidecar
|
||||
rather than a floating MP4 duration:
|
||||
|
||||
@@ -517,6 +561,10 @@ the following block. It duplicates the most recent available frame when the
|
||||
visual tower needs an even frame count. Audio is loaded once at 16 kHz and
|
||||
sliced at the same 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,17 +584,85 @@ python streaming_turn_model.py \
|
||||
--output features/Session_1_first_4s_check.pt
|
||||
```
|
||||
|
||||
<<<<<<< HEAD
|
||||
Do not launch all sessions unless `cache_replay` cosine is nearly 1. For
|
||||
`open-stream-replay`, a materially lower `continuous_window_semantics` cosine
|
||||
means that the custom payload stitching is not equivalent to one continuous
|
||||
Qwen input; it must not be dismissed as a media-turn-boundary effect because
|
||||
this mode intentionally keeps a single media turn open.
|
||||
=======
|
||||
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 blockwise media preprocessing changes the representation.
|
||||
>>>>>>> 8145be410866dd0675038d92347aea34d24b84d5
|
||||
|
||||
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
|
||||
same most-recent visual frame. A frame with timestamp later than `t` must never
|
||||
be included.
|
||||
|
||||
## Direct extraction without a prepared 2 FPS MP4
|
||||
|
||||
`streaming_turn_model_directly.py` removes the offline video-transcoding
|
||||
requirement. It reads original frame PTS with Decord, constructs a global
|
||||
2 FPS lattice, maps every lattice point to the latest source frame whose PTS
|
||||
is not later than that point, and passes explicit frame lists to Qwen. Audio is
|
||||
decoded once to 16-kHz mono and sliced at exact causal boundaries. Therefore
|
||||
no FFmpeg `fps` filter or second H.264/AAC encode can alter the experiment.
|
||||
|
||||
Validate one Ego4D clip without loading Qwen:
|
||||
|
||||
```bash
|
||||
python streaming_turn_model_directly.py \
|
||||
--video /data/wangx/Datasets/ego4d/v1/clips/CLIP_UID.mp4 \
|
||||
--target /data/wangx/Datasets/ego4d/ego4d_target_frame/CLIP_UID.npy \
|
||||
--video-fps 2 \
|
||||
--media-only
|
||||
```
|
||||
|
||||
Run a short native persistent-KV check:
|
||||
|
||||
```bash
|
||||
python streaming_turn_model_directly.py \
|
||||
--model-id /data/wangx/Models/Qwen2.5-Omni-3B \
|
||||
--video /data/wangx/Datasets/ego4d/v1/clips/CLIP_UID.mp4 \
|
||||
--end 40 \
|
||||
--mode open-stream-replay \
|
||||
--memory-seconds 0 \
|
||||
--video-fps 2 \
|
||||
--flash-attention-2 \
|
||||
--device cuda:0 \
|
||||
--output features/direct/CLIP_UID_first_40s.pt
|
||||
```
|
||||
|
||||
Here `--memory-seconds 0` means that completed-block Thinker KV persists for
|
||||
the entire requested interval. Do not use this setting blindly for complete
|
||||
multi-minute recordings: it has linear KV growth and remains subject to the
|
||||
Thinker's context limit. A full bounded-memory dataset run should explicitly
|
||||
use the experimental budget, for example `--memory-seconds 20`.
|
||||
|
||||
EasyCom can be read directly as the dense concatenation of retained clips;
|
||||
filename timecode gaps are ignored and the rightmost 1920x1080 region is kept:
|
||||
|
||||
```bash
|
||||
python streaming_turn_model_directly.py \
|
||||
--model-id /path/to/Qwen2.5-Omni-3B \
|
||||
--easycom-session-dir /path/to/Main/Video_Compressed/Session_1 \
|
||||
--target target_frame/Session_1.npy \
|
||||
--mode open-stream-replay \
|
||||
--memory-seconds 20 \
|
||||
--video-fps 2 \
|
||||
--flash-attention-2 \
|
||||
--device cuda:0 \
|
||||
--output features/direct/Session_1.pt
|
||||
```
|
||||
|
||||
Matching `states.shape[0]` to the target is necessary but not sufficient.
|
||||
The direct reader also validates EasyCom's 20-Hz source inventory, source PTS
|
||||
monotonicity, dense duration, crop bounds, audio sample count, and causal
|
||||
frame-to-grid mapping. The original prepare scripts remain useful only as an
|
||||
optional decode-speed/reproducibility cache.
|
||||
|
||||
## Saved streaming tensors
|
||||
|
||||
Each `.pt` file contains:
|
||||
@@ -556,7 +672,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;
|
||||
@@ -580,7 +696,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 (
|
||||
@@ -589,16 +706,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/open_2fps \
|
||||
--target-root target_frame \
|
||||
--output checkpoints/open_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.
|
||||
@@ -610,7 +748,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/open_2fps \
|
||||
--target-root target_frame \
|
||||
--checkpoint checkpoints/open_2fps_head.pth \
|
||||
--output-dir results/open_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]
|
||||
@@ -626,8 +779,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:
|
||||
|
||||
@@ -706,7 +873,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`.
|
||||
@@ -737,8 +904,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;
|
||||
@@ -751,7 +919,8 @@ Once these checks pass, the first publishable experiment is:
|
||||
```text
|
||||
frozen Qwen2.5-Omni-3B Thinker
|
||||
+ trained 10x3 head
|
||||
+ 20-second causal memory
|
||||
+ bounded Qwen-style 2-second open streaming
|
||||
+ 20-second rolling causal memory
|
||||
+ full-session 5 Hz output
|
||||
+ exact EgoSpeak split and Table 2 metric
|
||||
```
|
||||
|
||||
299
easycom_head.py
Normal file
299
easycom_head.py
Normal file
@@ -0,0 +1,299 @@
|
||||
"""Shared head, labels, and metrics for the EasyCom EgoSpeak experiment."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any, Literal, Sequence
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
OFFICIAL_TRAIN_SESSIONS = (
|
||||
"Session_7",
|
||||
"Session_11",
|
||||
"Session_10",
|
||||
"Session_6",
|
||||
"Session_9",
|
||||
"Session_4",
|
||||
"Session_8",
|
||||
"Session_5",
|
||||
"Session_12",
|
||||
)
|
||||
OFFICIAL_TEST_SESSIONS = ("Session_3", "Session_2", "Session_1")
|
||||
CLASS_NAMES = (
|
||||
"background",
|
||||
"camera_wearer_speaking",
|
||||
"other_people_speaking",
|
||||
)
|
||||
STEP_SECONDS = 0.2
|
||||
NUM_HORIZONS = 10
|
||||
NUM_CLASSES = 3
|
||||
|
||||
|
||||
class TurnPredictionHead(nn.Module):
|
||||
"""Map one Omni state to EgoSpeak-compatible ``[10, 3]`` logits."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int = 2048,
|
||||
num_horizons: int = NUM_HORIZONS,
|
||||
num_classes: int = NUM_CLASSES,
|
||||
dropout: float = 0.1,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.hidden_size = hidden_size
|
||||
self.num_horizons = num_horizons
|
||||
self.num_classes = num_classes
|
||||
self.dropout = dropout
|
||||
self.net = nn.Sequential(
|
||||
nn.LayerNorm(hidden_size),
|
||||
nn.Linear(hidden_size, hidden_size // 2),
|
||||
nn.GELU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_size // 2, num_horizons * num_classes),
|
||||
)
|
||||
|
||||
def forward(self, states: torch.Tensor) -> torch.Tensor:
|
||||
logits = self.net(states)
|
||||
return logits.view(
|
||||
*states.shape[:-1], self.num_horizons, self.num_classes
|
||||
)
|
||||
|
||||
|
||||
def _validate_loss_inputs(logits: torch.Tensor, targets: torch.Tensor) -> None:
|
||||
if logits.shape != targets.shape:
|
||||
raise ValueError(
|
||||
f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
|
||||
f"{tuple(targets.shape)}"
|
||||
)
|
||||
if logits.shape[-1] != NUM_CLASSES:
|
||||
raise ValueError(
|
||||
f"Expected {NUM_CLASSES} classes, got {logits.shape[-1]}"
|
||||
)
|
||||
|
||||
|
||||
def egospeak_multihorizon_loss(
|
||||
logits: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
*,
|
||||
reduction: Literal["mean", "sum", "none"] = "mean",
|
||||
) -> torch.Tensor:
|
||||
"""Reproduce the LSTR ``MCE`` loss used by the supplied shell/config.
|
||||
|
||||
EasyCom targets may contain both active speaker classes as ``[0, 1, 1]``.
|
||||
LSTR preserves those weights directly and sums their negative log-softmax
|
||||
terms. It does not normalize the target row.
|
||||
"""
|
||||
|
||||
_validate_loss_inputs(logits, targets)
|
||||
targets = targets.to(device=logits.device, dtype=logits.dtype)
|
||||
loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
|
||||
if reduction == "mean":
|
||||
return loss.mean()
|
||||
if reduction == "sum":
|
||||
return loss.sum()
|
||||
if reduction == "none":
|
||||
return loss
|
||||
raise ValueError(f"Unknown reduction: {reduction}")
|
||||
|
||||
|
||||
def miniroad_multihorizon_loss(
|
||||
logits: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
*,
|
||||
reduction: Literal["mean", "sum", "none"] = "mean",
|
||||
) -> torch.Tensor:
|
||||
"""Reproduce MiniROAD's L2-normalized multi-positive anticipation loss."""
|
||||
|
||||
_validate_loss_inputs(logits, targets)
|
||||
targets = nn.functional.normalize(
|
||||
targets.to(device=logits.device, dtype=logits.dtype),
|
||||
p=2,
|
||||
dim=-1,
|
||||
)
|
||||
loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
|
||||
if reduction == "mean":
|
||||
return loss.mean()
|
||||
if reduction == "sum":
|
||||
return loss.sum()
|
||||
if reduction == "none":
|
||||
return loss
|
||||
raise ValueError(f"Unknown reduction: {reduction}")
|
||||
|
||||
|
||||
def multihorizon_cross_entropy(
|
||||
logits: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Backward-compatible alias for the LSTR/EgoSpeak loss."""
|
||||
|
||||
return egospeak_multihorizon_loss(logits, targets)
|
||||
|
||||
|
||||
def normalize_session_name(value: str | int) -> str:
|
||||
text = str(value)
|
||||
if text.startswith("Session_"):
|
||||
suffix = text.removeprefix("Session_")
|
||||
else:
|
||||
suffix = text
|
||||
if not suffix.isdigit():
|
||||
raise ValueError(f"Invalid EasyCom session: {value}")
|
||||
return f"Session_{int(suffix)}"
|
||||
|
||||
|
||||
@dataclass
|
||||
class SessionData:
|
||||
name: str
|
||||
states: torch.Tensor
|
||||
targets: torch.Tensor
|
||||
label_indices: torch.LongTensor
|
||||
feature_metadata: dict[str, Any]
|
||||
|
||||
@property
|
||||
def hidden_size(self) -> int:
|
||||
return int(self.states.shape[-1])
|
||||
|
||||
def future_examples(
|
||||
self,
|
||||
num_horizons: int = NUM_HORIZONS,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, torch.LongTensor]:
|
||||
"""Return states, future labels, and target indices for valid rows."""
|
||||
|
||||
offsets = torch.arange(1, num_horizons + 1, dtype=torch.long)
|
||||
future_indices = self.label_indices[:, None] + offsets[None, :]
|
||||
valid = (future_indices >= 0).all(dim=1) & (
|
||||
future_indices < len(self.targets)
|
||||
).all(dim=1)
|
||||
rows = future_indices[valid]
|
||||
return self.states[valid], self.targets[rows], rows
|
||||
|
||||
|
||||
def load_session_data(
|
||||
feature_root: str | Path,
|
||||
target_root: str | Path,
|
||||
session: str | int,
|
||||
*,
|
||||
expected_mode: str | None = "causal-prefix",
|
||||
) -> SessionData:
|
||||
"""Load and validate one saved Omni feature payload and target array."""
|
||||
|
||||
name = normalize_session_name(session)
|
||||
feature_path = Path(feature_root) / f"{name}.pt"
|
||||
target_path = Path(target_root) / f"{name}.npy"
|
||||
if not feature_path.is_file():
|
||||
raise FileNotFoundError(feature_path)
|
||||
if not target_path.is_file():
|
||||
raise FileNotFoundError(target_path)
|
||||
|
||||
payload = torch.load(feature_path, map_location="cpu", weights_only=False)
|
||||
if not isinstance(payload, dict) or "states" not in payload:
|
||||
raise ValueError(f"Invalid feature payload: {feature_path}")
|
||||
states = payload["states"]
|
||||
if not isinstance(states, torch.Tensor) or states.ndim != 2:
|
||||
raise ValueError(f"states must be [T,D] in {feature_path}")
|
||||
if not torch.is_floating_point(states):
|
||||
raise ValueError(f"states must be floating point in {feature_path}")
|
||||
if not torch.isfinite(states).all():
|
||||
raise ValueError(f"states contain NaN/Inf in {feature_path}")
|
||||
|
||||
mode = payload.get("mode")
|
||||
if expected_mode is not None and mode != expected_mode:
|
||||
raise ValueError(
|
||||
f"{feature_path} mode is {mode!r}, expected {expected_mode!r}"
|
||||
)
|
||||
label_indices = payload.get("label_indices")
|
||||
if label_indices is None:
|
||||
raise ValueError(f"Missing label_indices in {feature_path}")
|
||||
label_indices = torch.as_tensor(label_indices, dtype=torch.long)
|
||||
if label_indices.ndim != 1 or len(label_indices) != len(states):
|
||||
raise ValueError(f"Invalid label_indices in {feature_path}")
|
||||
if len(label_indices) and not torch.equal(
|
||||
label_indices,
|
||||
torch.arange(len(label_indices), dtype=torch.long),
|
||||
):
|
||||
raise ValueError(
|
||||
f"{feature_path} does not use the verified zero-based dense alignment"
|
||||
)
|
||||
|
||||
target_array = np.load(target_path, allow_pickle=False)
|
||||
if target_array.ndim != 2 or target_array.shape[1] != NUM_CLASSES:
|
||||
raise ValueError(
|
||||
f"Expected target [T,{NUM_CLASSES}], got {target_array.shape}: "
|
||||
f"{target_path}"
|
||||
)
|
||||
targets = torch.from_numpy(np.asarray(target_array, dtype=np.float32))
|
||||
if len(label_indices) and int(label_indices[-1]) >= len(targets):
|
||||
raise ValueError(
|
||||
f"Last label index {int(label_indices[-1])} exceeds target length "
|
||||
f"{len(targets)} in {name}"
|
||||
)
|
||||
if len(states) != len(targets):
|
||||
raise ValueError(
|
||||
f"{name}: feature/target length mismatch: {len(states)} vs "
|
||||
f"{len(targets)}"
|
||||
)
|
||||
|
||||
metadata = {
|
||||
key: payload.get(key)
|
||||
for key in (
|
||||
"model_id",
|
||||
"mode",
|
||||
"config",
|
||||
"start",
|
||||
"end",
|
||||
"num_steps",
|
||||
"video_sampling_grid",
|
||||
)
|
||||
}
|
||||
return SessionData(name, states, targets, label_indices, metadata)
|
||||
|
||||
|
||||
def average_precision_metrics(
|
||||
scores: np.ndarray,
|
||||
targets: np.ndarray,
|
||||
) -> dict[str, Any]:
|
||||
"""Compute EgoSpeak frame AP for all three classes, including background."""
|
||||
|
||||
from sklearn.metrics import average_precision_score
|
||||
|
||||
scores = np.asarray(scores, dtype=np.float64)
|
||||
targets = np.asarray(targets, dtype=np.float64)
|
||||
if scores.shape != targets.shape or scores.ndim != 2:
|
||||
raise ValueError(
|
||||
f"scores/targets must have equal [N,C] shape, got "
|
||||
f"{scores.shape} and {targets.shape}"
|
||||
)
|
||||
if scores.shape[1] != NUM_CLASSES:
|
||||
raise ValueError(f"Expected {NUM_CLASSES} classes, got {scores.shape[1]}")
|
||||
|
||||
aps = average_precision_score(targets, scores, average=None)
|
||||
return {
|
||||
"num_steps": int(len(scores)),
|
||||
"class_ap": {
|
||||
name: float(value) for name, value in zip(CLASS_NAMES, aps)
|
||||
},
|
||||
"map": float(np.mean(aps)),
|
||||
}
|
||||
|
||||
|
||||
def checkpoint_head_config(payload: dict[str, Any]) -> dict[str, Any]:
|
||||
"""Return a validated head config, including legacy-checkpoint defaults."""
|
||||
|
||||
config = dict(payload.get("head_config", {}))
|
||||
config.setdefault("hidden_size", 2048)
|
||||
config.setdefault("num_horizons", NUM_HORIZONS)
|
||||
config.setdefault("num_classes", NUM_CLASSES)
|
||||
config.setdefault("dropout", 0.1)
|
||||
return config
|
||||
|
||||
|
||||
def validate_disjoint_sessions(
|
||||
first: Sequence[str],
|
||||
second: Sequence[str],
|
||||
) -> None:
|
||||
overlap = set(first) & set(second)
|
||||
if overlap:
|
||||
raise ValueError(f"Session sets overlap: {sorted(overlap)}")
|
||||
289
infer_turn_head.py
Normal file
289
infer_turn_head.py
Normal file
@@ -0,0 +1,289 @@
|
||||
"""Run the Omni anticipation head and evaluate it with EgoSpeak's EasyCom metric."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from easycom_head import (
|
||||
CLASS_NAMES,
|
||||
NUM_CLASSES,
|
||||
NUM_HORIZONS,
|
||||
OFFICIAL_TEST_SESSIONS,
|
||||
STEP_SECONDS,
|
||||
TurnPredictionHead,
|
||||
average_precision_metrics,
|
||||
checkpoint_head_config,
|
||||
load_session_data,
|
||||
normalize_session_name,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--feature-root", type=Path, required=True)
|
||||
parser.add_argument("--target-root", type=Path, required=True)
|
||||
parser.add_argument("--checkpoint", type=Path, required=True)
|
||||
parser.add_argument("--output-dir", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--sessions",
|
||||
nargs="+",
|
||||
default=list(OFFICIAL_TEST_SESSIONS),
|
||||
help="Evaluation sessions; default is the official Session 3/2/1 test split.",
|
||||
)
|
||||
parser.add_argument("--batch-size", type=int, default=1024)
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument(
|
||||
"--expected-mode",
|
||||
default="causal-prefix",
|
||||
help="Expected feature mode, or 'any' to disable the check.",
|
||||
)
|
||||
parser.add_argument("--no-progress", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _resolve_device(name: str) -> torch.device:
|
||||
device = torch.device(name)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(f"CUDA requested but unavailable: {name}")
|
||||
return device
|
||||
|
||||
|
||||
def _load_head(
|
||||
checkpoint: Path,
|
||||
device: torch.device,
|
||||
) -> tuple[TurnPredictionHead, dict[str, Any]]:
|
||||
payload = torch.load(checkpoint, map_location="cpu", weights_only=False)
|
||||
if not isinstance(payload, dict):
|
||||
raise ValueError(f"Invalid checkpoint: {checkpoint}")
|
||||
state_dict = payload.get("head")
|
||||
if not isinstance(state_dict, dict):
|
||||
raise ValueError(f"Checkpoint has no head state_dict: {checkpoint}")
|
||||
config = checkpoint_head_config(payload)
|
||||
if int(config["num_horizons"]) != NUM_HORIZONS:
|
||||
raise ValueError(
|
||||
f"Checkpoint has {config['num_horizons']} horizons, expected {NUM_HORIZONS}"
|
||||
)
|
||||
if int(config["num_classes"]) != NUM_CLASSES:
|
||||
raise ValueError(
|
||||
f"Checkpoint has {config['num_classes']} classes, expected {NUM_CLASSES}"
|
||||
)
|
||||
head = TurnPredictionHead(**config)
|
||||
head.load_state_dict(state_dict)
|
||||
return head.to(device).eval(), payload
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _predict_states(
|
||||
head: TurnPredictionHead,
|
||||
states: torch.Tensor,
|
||||
device: torch.device,
|
||||
batch_size: int,
|
||||
*,
|
||||
description: str,
|
||||
show_progress: bool,
|
||||
) -> np.ndarray:
|
||||
chunks: list[torch.Tensor] = []
|
||||
starts = range(0, len(states), batch_size)
|
||||
progress = tqdm(
|
||||
starts,
|
||||
total=(len(states) + batch_size - 1) // batch_size,
|
||||
desc=description,
|
||||
unit="batch",
|
||||
disable=not show_progress,
|
||||
)
|
||||
for start in progress:
|
||||
batch = states[start : start + batch_size].to(
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
non_blocking=True,
|
||||
)
|
||||
chunks.append(head(batch).softmax(dim=-1).cpu())
|
||||
return torch.cat(chunks, dim=0).numpy()
|
||||
|
||||
|
||||
def _align_to_egospeak_timeline(
|
||||
probabilities: np.ndarray,
|
||||
label_indices: torch.LongTensor,
|
||||
num_target_steps: int,
|
||||
) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Return EgoSpeak ``[T,3,10]`` scores and their causal-valid mask."""
|
||||
|
||||
predictions = np.zeros(
|
||||
(num_target_steps, NUM_CLASSES, NUM_HORIZONS),
|
||||
dtype=np.float32,
|
||||
)
|
||||
valid = np.zeros((num_target_steps, NUM_HORIZONS), dtype=bool)
|
||||
current = label_indices.numpy()
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
target_indices = current + horizon + 1
|
||||
usable = (target_indices >= 0) & (target_indices < num_target_steps)
|
||||
rows = target_indices[usable]
|
||||
predictions[rows, :, horizon] = probabilities[usable, horizon, :]
|
||||
valid[rows, horizon] = True
|
||||
return predictions, valid
|
||||
|
||||
|
||||
def _evaluate(
|
||||
session_predictions: dict[str, np.ndarray],
|
||||
session_valid: dict[str, np.ndarray],
|
||||
session_targets: dict[str, np.ndarray],
|
||||
) -> dict[str, Any]:
|
||||
full_results = []
|
||||
valid_results = []
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
full_scores = np.concatenate(
|
||||
[scores[:, :, horizon] for scores in session_predictions.values()],
|
||||
axis=0,
|
||||
)
|
||||
full_targets = np.concatenate(list(session_targets.values()), axis=0)
|
||||
full = average_precision_metrics(full_scores, full_targets)
|
||||
full["horizon"] = horizon + 1
|
||||
full["seconds"] = (horizon + 1) * STEP_SECONDS
|
||||
full_results.append(full)
|
||||
|
||||
valid_scores = np.concatenate(
|
||||
[
|
||||
session_predictions[name][session_valid[name][:, horizon], :, horizon]
|
||||
for name in session_predictions
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
valid_targets = np.concatenate(
|
||||
[
|
||||
session_targets[name][session_valid[name][:, horizon]]
|
||||
for name in session_predictions
|
||||
],
|
||||
axis=0,
|
||||
)
|
||||
causal = average_precision_metrics(valid_scores, valid_targets)
|
||||
causal["horizon"] = horizon + 1
|
||||
causal["seconds"] = (horizon + 1) * STEP_SECONDS
|
||||
valid_results.append(causal)
|
||||
|
||||
return {
|
||||
"egospeak_full_timeline": {
|
||||
"description": (
|
||||
"Exact EgoSpeak array/evaluator layout. Unavailable pre-session "
|
||||
"causal predictions remain zero in the preallocated arrays."
|
||||
),
|
||||
"horizons": full_results,
|
||||
"mean_map": float(np.mean([item["map"] for item in full_results])),
|
||||
},
|
||||
"causal_valid_only": {
|
||||
"description": (
|
||||
"Primary causal metric; excludes only the first h rows of each "
|
||||
"session for horizon h because no pre-session context exists."
|
||||
),
|
||||
"horizons": valid_results,
|
||||
"mean_map": float(np.mean([item["map"] for item in valid_results])),
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _print_metrics(metrics: dict[str, Any]) -> None:
|
||||
for protocol_name in ("egospeak_full_timeline", "causal_valid_only"):
|
||||
protocol = metrics[protocol_name]
|
||||
print(f"\n{protocol_name}:")
|
||||
for result in protocol["horizons"]:
|
||||
print(
|
||||
f"Action anticipation ({result['seconds']:.2f}s) "
|
||||
f"perframe mAP: {result['map'] * 100:.5f}"
|
||||
)
|
||||
class_text = ", ".join(
|
||||
f"{name}={value * 100:.5f}"
|
||||
for name, value in result["class_ap"].items()
|
||||
)
|
||||
print(f" AP: {class_text}; n={result['num_steps']}")
|
||||
print(
|
||||
f"Action anticipation (mean) perframe mAP: "
|
||||
f"{protocol['mean_map'] * 100:.5f}"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
if args.batch_size <= 0:
|
||||
raise ValueError("--batch-size must be positive")
|
||||
device = _resolve_device(args.device)
|
||||
expected_mode = None if args.expected_mode == "any" else args.expected_mode
|
||||
sessions = [normalize_session_name(item) for item in args.sessions]
|
||||
if len(sessions) != len(set(sessions)):
|
||||
raise ValueError("Duplicate sessions were supplied")
|
||||
head, checkpoint_payload = _load_head(args.checkpoint, device)
|
||||
head_config = checkpoint_head_config(checkpoint_payload)
|
||||
|
||||
args.output_dir.mkdir(parents=True, exist_ok=True)
|
||||
prediction_dir = args.output_dir / "pred"
|
||||
valid_dir = args.output_dir / "valid"
|
||||
prediction_dir.mkdir(parents=True, exist_ok=True)
|
||||
valid_dir.mkdir(parents=True, exist_ok=True)
|
||||
session_predictions: dict[str, np.ndarray] = {}
|
||||
session_valid: dict[str, np.ndarray] = {}
|
||||
session_targets: dict[str, np.ndarray] = {}
|
||||
|
||||
print(f"device: {device}")
|
||||
print(f"sessions: {', '.join(sessions)}")
|
||||
for name in sessions:
|
||||
data = load_session_data(
|
||||
args.feature_root,
|
||||
args.target_root,
|
||||
name,
|
||||
expected_mode=expected_mode,
|
||||
)
|
||||
if data.hidden_size != int(head_config["hidden_size"]):
|
||||
raise ValueError(
|
||||
f"{name} feature dim is {data.hidden_size}, checkpoint expects "
|
||||
f"{head_config['hidden_size']}"
|
||||
)
|
||||
probabilities = _predict_states(
|
||||
head,
|
||||
data.states,
|
||||
device,
|
||||
args.batch_size,
|
||||
description=f"infer {name}",
|
||||
show_progress=not args.no_progress,
|
||||
)
|
||||
predictions, valid = _align_to_egospeak_timeline(
|
||||
probabilities,
|
||||
data.label_indices,
|
||||
len(data.targets),
|
||||
)
|
||||
target_array = data.targets.numpy()
|
||||
session_predictions[name] = predictions
|
||||
session_valid[name] = valid
|
||||
session_targets[name] = target_array
|
||||
np.save(prediction_dir / f"{name}.npy", predictions)
|
||||
np.save(valid_dir / f"{name}.npy", valid)
|
||||
print(
|
||||
f"{name}: saved pred {predictions.shape}, "
|
||||
f"causal-valid counts={valid.sum(axis=0).tolist()}"
|
||||
)
|
||||
|
||||
metrics = _evaluate(session_predictions, session_valid, session_targets)
|
||||
metrics.update(
|
||||
{
|
||||
"sessions": sessions,
|
||||
"classes": list(CLASS_NAMES),
|
||||
"prediction_layout": "[target_time, class, future_horizon]",
|
||||
"checkpoint": str(args.checkpoint.resolve()),
|
||||
"checkpoint_epoch": checkpoint_payload.get("epoch"),
|
||||
"checkpoint_train_sessions": checkpoint_payload.get("train_sessions"),
|
||||
"checkpoint_dev_sessions": checkpoint_payload.get("dev_sessions"),
|
||||
}
|
||||
)
|
||||
metrics_path = args.output_dir / "metrics.json"
|
||||
metrics_path.write_text(json.dumps(metrics, indent=2) + "\n")
|
||||
_print_metrics(metrics)
|
||||
print(f"\nmetrics: {metrics_path}")
|
||||
print(f"EgoSpeak-compatible predictions: {prediction_dir}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
@@ -42,10 +42,22 @@ class ClipInfo:
|
||||
|
||||
|
||||
def _run_json(command: list[str]) -> dict:
|
||||
completed = subprocess.run(
|
||||
command, check=True, capture_output=True, text=True
|
||||
)
|
||||
return json.loads(completed.stdout)
|
||||
try:
|
||||
completed = subprocess.run(
|
||||
command, check=True, capture_output=True, text=True
|
||||
)
|
||||
except subprocess.CalledProcessError as exc:
|
||||
stderr = (exc.stderr or "").strip()
|
||||
detail = f"\n{stderr}" if stderr else ""
|
||||
raise RuntimeError(
|
||||
f"Command failed ({exc.returncode}): {shlex.join(command)}{detail}"
|
||||
) from exc
|
||||
try:
|
||||
return json.loads(completed.stdout)
|
||||
except json.JSONDecodeError as exc:
|
||||
raise RuntimeError(
|
||||
f"Command returned invalid JSON: {shlex.join(command)}"
|
||||
) from exc
|
||||
|
||||
|
||||
def probe_clip(path: Path) -> ClipInfo:
|
||||
@@ -108,16 +120,22 @@ def validate_inventory(clips: list[ClipInfo], target_steps: int | None) -> int:
|
||||
f"{clip.width}x{clip.height}: {clip.path}"
|
||||
)
|
||||
total_frames = sum(clip.frames for clip in clips)
|
||||
if total_frames % (SOURCE_FPS // TASK_FPS):
|
||||
raise ValueError(
|
||||
f"Total source frames {total_frames} are not divisible into 5 Hz steps"
|
||||
)
|
||||
media_steps = total_frames // (SOURCE_FPS // TASK_FPS)
|
||||
frames_per_step = SOURCE_FPS // TASK_FPS
|
||||
media_steps = total_frames // frames_per_step
|
||||
if target_steps is not None and media_steps != target_steps:
|
||||
raise ValueError(
|
||||
f"Dense media gives {media_steps} steps but target has {target_steps}"
|
||||
f"Dense media gives {media_steps} complete 5 Hz steps from "
|
||||
f"{total_frames} frames, but target has {target_steps}"
|
||||
)
|
||||
return media_steps
|
||||
# A session may end with 1--3 source frames, which cannot form another
|
||||
# complete 5 Hz target interval. The target length is authoritative; the
|
||||
# final video/audio trim below discards only this sub-step tail.
|
||||
return target_steps if target_steps is not None else media_steps
|
||||
|
||||
|
||||
def fixed_grid_frame_count(*, fps: float, dense_duration: float) -> int:
|
||||
"""Number of samples on the grid ``n / fps < dense_duration``."""
|
||||
return int(math.ceil(dense_duration * fps - 1e-7))
|
||||
|
||||
|
||||
def build_ffmpeg_command(
|
||||
@@ -127,6 +145,9 @@ def build_ffmpeg_command(
|
||||
fps: float,
|
||||
dense_duration: float,
|
||||
) -> list[str]:
|
||||
expected_frames = fixed_grid_frame_count(
|
||||
fps=fps, dense_duration=dense_duration
|
||||
)
|
||||
# Trim every decoded stream independently before concatenation. Released
|
||||
# AAC clips can contain roughly one codec-padding frame beyond their video
|
||||
# duration; concatenating them first would accumulate AV drift.
|
||||
@@ -158,11 +179,19 @@ def build_ffmpeg_command(
|
||||
"".join(concat_inputs)
|
||||
+ f"concat=n={len(clips)}:v=1:a=1[vcat][acat]"
|
||||
)
|
||||
# fps=start_time=0 establishes one global visual lattice after the clips
|
||||
# have become a dense timeline. The final atrim is a defensive guard
|
||||
# against decoder rounding at the last sample.
|
||||
# Establish one global visual lattice after the clips have become a dense
|
||||
# timeline. Some FFmpeg versions omit the last grid point when the input
|
||||
# ends exactly on an integer-second boundary. Add two disposable clones,
|
||||
# then trim by frame index, so EOF timestamp rounding cannot change the
|
||||
# number of samples. The clones are retained only if a real final grid
|
||||
# sample was missing.
|
||||
filters.append(
|
||||
f"[vcat]fps=fps={fps}:start_time=0:round=down,format=yuv420p[vout]"
|
||||
f"[vcat]trim=start=0:end={dense_duration:.9f},"
|
||||
f"setpts=PTS-STARTPTS,"
|
||||
f"fps=fps={fps}:start_time=0:round=down,"
|
||||
f"tpad=stop_mode=clone:stop=2,"
|
||||
f"trim=start_frame=0:end_frame={expected_frames},"
|
||||
f"setpts=N/({fps}*TB),format=yuv420p[vout]"
|
||||
)
|
||||
filters.append(
|
||||
f"[acat]atrim=start=0:end={dense_duration:.9f},"
|
||||
@@ -220,20 +249,15 @@ def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict:
|
||||
raise RuntimeError(f"Prepared file must contain video and audio: {path}")
|
||||
actual_fps = float(Fraction(video["avg_frame_rate"]))
|
||||
actual_frames = int(video["nb_frames"])
|
||||
# FFmpeg may either retain or discard the final partial visual interval,
|
||||
# depending on timestamps delivered by the concat filter. Both are valid
|
||||
# fixed grids; for 1388.4 s at 2 FPS this permits 2776 or 2777 frames.
|
||||
frame_product = dense_duration * fps
|
||||
expected_frames = {
|
||||
int(math.floor(frame_product + 1e-7)),
|
||||
int(math.ceil(frame_product - 1e-7)),
|
||||
}
|
||||
expected_frames = fixed_grid_frame_count(
|
||||
fps=fps, dense_duration=dense_duration
|
||||
)
|
||||
if not math.isclose(actual_fps, fps, rel_tol=0.0, abs_tol=1e-6):
|
||||
raise RuntimeError(f"Prepared FPS is {actual_fps}, expected {fps}: {path}")
|
||||
if actual_frames not in expected_frames:
|
||||
if actual_frames != expected_frames:
|
||||
raise RuntimeError(
|
||||
f"Prepared frame count is {actual_frames}, expected one of "
|
||||
f"{sorted(expected_frames)}: {path}"
|
||||
f"Prepared frame count is {actual_frames}, expected "
|
||||
f"{expected_frames}: {path}"
|
||||
)
|
||||
audio_duration = float(audio["duration"])
|
||||
# AAC frame padding can shift the container duration by a few milliseconds.
|
||||
@@ -245,7 +269,7 @@ def validate_output(path: Path, *, fps: float, dense_duration: float) -> dict:
|
||||
return {
|
||||
"video_fps": actual_fps,
|
||||
"video_frames": actual_frames,
|
||||
"allowed_video_frames": sorted(expected_frames),
|
||||
"expected_video_frames": expected_frames,
|
||||
"video_duration_seconds": float(video["duration"]),
|
||||
"audio_duration_seconds": audio_duration,
|
||||
"audio_sample_rate": int(audio["sample_rate"]),
|
||||
@@ -322,12 +346,16 @@ def main() -> None:
|
||||
output, fps=fps, dense_duration=duration
|
||||
)
|
||||
|
||||
source_frames = sum(clip.frames for clip in clips)
|
||||
used_source_frames = num_steps * (SOURCE_FPS // TASK_FPS)
|
||||
metadata = {
|
||||
"session": session,
|
||||
"source_session_dir": str(args.session_dir.resolve()),
|
||||
"source_clip_order": [asdict(clip) for clip in clips],
|
||||
"source_num_clips": len(clips),
|
||||
"source_frames_20hz": sum(clip.frames for clip in clips),
|
||||
"source_frames_20hz": source_frames,
|
||||
"used_source_frames_20hz": used_source_frames,
|
||||
"discarded_tail_frames_20hz": source_frames - used_source_frames,
|
||||
"task_fps": TASK_FPS,
|
||||
"num_steps": num_steps,
|
||||
"dense_duration_seconds": duration,
|
||||
|
||||
421
prepare_ego4d_2fps.py
Executable file
421
prepare_ego4d_2fps.py
Executable file
@@ -0,0 +1,421 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Prepare EgoSpeak's Ego4D split as exact 2 FPS audio-video clips.
|
||||
|
||||
The 5 Hz target arrays are the timing authority. For a target with T rows,
|
||||
the prepared clip has duration T / 5 seconds and exactly ceil(2T / 5) video
|
||||
frames. Short source audio is padded with silence and all audio is trimmed to
|
||||
the target duration.
|
||||
|
||||
Processing is intentionally sequential. A single libx264 process uses all
|
||||
available CPU cores on the intended extraction machine, so launching several
|
||||
encoders normally creates CPU and disk contention instead of improving total
|
||||
throughput.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import json
|
||||
import math
|
||||
import shutil
|
||||
import subprocess
|
||||
import sys
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
|
||||
import numpy as np
|
||||
|
||||
|
||||
TARGET_FPS = 5
|
||||
VIDEO_FPS = 2
|
||||
EXPECTED_TRAIN_CLIPS = 346
|
||||
EXPECTED_TEST_CLIPS = 87
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class ClipSpec:
|
||||
split: str
|
||||
uid: str
|
||||
steps: int
|
||||
duration: float
|
||||
video_frames: int
|
||||
source: Path
|
||||
target: Path
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(
|
||||
description=(
|
||||
"Prepare the 433 clips in EgoSpeak's EGO4DSHUFFLE split as "
|
||||
"strict fixed-grid 2 FPS MP4 files."
|
||||
)
|
||||
)
|
||||
parser.add_argument(
|
||||
"--data-root",
|
||||
type=Path,
|
||||
default=Path("/data/wangx/Datasets/ego4d"),
|
||||
help="Ego4D dataset root (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split-json",
|
||||
type=Path,
|
||||
help="EgoSpeak LSTR data_info.json; defaults to DATA_ROOT/egospeak_data_info.json",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--video-root",
|
||||
type=Path,
|
||||
help="Raw clip directory; defaults to DATA_ROOT/v1/clips",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target-root",
|
||||
type=Path,
|
||||
help="Target directory; defaults to DATA_ROOT/ego4d_target_frame",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--output-root",
|
||||
type=Path,
|
||||
help="Prepared video directory; defaults to DATA_ROOT/omni_prepared_2fps",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--manifest",
|
||||
type=Path,
|
||||
help="Output TSV; defaults to DATA_ROOT/ego4d_egospeak.tsv",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--split",
|
||||
choices=("all", "train", "test"),
|
||||
default="all",
|
||||
help="Prepare all, train, or test clips (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--uid",
|
||||
action="append",
|
||||
help="Prepare only this split UID; repeat to select several clips",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--limit",
|
||||
type=int,
|
||||
help="Prepare only the first N selected clips, for testing",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--crf",
|
||||
type=int,
|
||||
default=18,
|
||||
help="libx264 constant-rate-factor value (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--preset",
|
||||
default="medium",
|
||||
help="libx264 preset (default: %(default)s)",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--force",
|
||||
action="store_true",
|
||||
help="Rebuild outputs even when their frame counts are already correct",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dry-run",
|
||||
action="store_true",
|
||||
help="Validate inputs and print work without running ffmpeg",
|
||||
)
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.limit is not None and args.limit <= 0:
|
||||
parser.error("--limit must be positive")
|
||||
if not 0 <= args.crf <= 51:
|
||||
parser.error("--crf must be between 0 and 51")
|
||||
return args
|
||||
|
||||
|
||||
def resolve_paths(args: argparse.Namespace) -> None:
|
||||
args.data_root = args.data_root.expanduser().resolve()
|
||||
args.split_json = (
|
||||
args.split_json or args.data_root / "egospeak_data_info.json"
|
||||
).expanduser().resolve()
|
||||
args.video_root = (
|
||||
args.video_root or args.data_root / "v1" / "clips"
|
||||
).expanduser().resolve()
|
||||
args.target_root = (
|
||||
args.target_root or args.data_root / "ego4d_target_frame"
|
||||
).expanduser().resolve()
|
||||
args.output_root = (
|
||||
args.output_root or args.data_root / "omni_prepared_2fps"
|
||||
).expanduser().resolve()
|
||||
args.manifest = (
|
||||
args.manifest or args.data_root / "ego4d_egospeak.tsv"
|
||||
).expanduser().resolve()
|
||||
|
||||
|
||||
def require_file(path: Path, description: str) -> None:
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"Missing {description}: {path}")
|
||||
|
||||
|
||||
def load_specs(args: argparse.Namespace) -> list[ClipSpec]:
|
||||
require_file(args.split_json, "EgoSpeak split JSON")
|
||||
if not args.video_root.is_dir():
|
||||
raise FileNotFoundError(f"Missing raw-video directory: {args.video_root}")
|
||||
if not args.target_root.is_dir():
|
||||
raise FileNotFoundError(f"Missing target directory: {args.target_root}")
|
||||
|
||||
with args.split_json.open() as file:
|
||||
payload = json.load(file)
|
||||
try:
|
||||
split_info = payload["EGO4DSHUFFLE"]
|
||||
train_uids = split_info["train_session_set"]
|
||||
test_uids = split_info["test_session_set"]
|
||||
except KeyError as error:
|
||||
raise RuntimeError(
|
||||
f"{args.split_json} is not EgoSpeak's LSTR data_info.json: "
|
||||
f"missing {error}"
|
||||
) from error
|
||||
|
||||
if len(train_uids) != EXPECTED_TRAIN_CLIPS:
|
||||
raise RuntimeError(
|
||||
f"Expected {EXPECTED_TRAIN_CLIPS} EGO4DSHUFFLE train clips, "
|
||||
f"found {len(train_uids)}"
|
||||
)
|
||||
if len(test_uids) != EXPECTED_TEST_CLIPS:
|
||||
raise RuntimeError(
|
||||
f"Expected {EXPECTED_TEST_CLIPS} EGO4DSHUFFLE test clips, "
|
||||
f"found {len(test_uids)}"
|
||||
)
|
||||
|
||||
specs: list[ClipSpec] = []
|
||||
seen: set[str] = set()
|
||||
for split, uids in (("train", train_uids), ("test", test_uids)):
|
||||
for uid in uids:
|
||||
if uid in seen:
|
||||
raise RuntimeError(f"Duplicate UID in EGO4DSHUFFLE: {uid}")
|
||||
seen.add(uid)
|
||||
|
||||
source = args.video_root / f"{uid}.mp4"
|
||||
target_path = args.target_root / f"{uid}.npy"
|
||||
require_file(source, "source video")
|
||||
require_file(target_path, "target array")
|
||||
|
||||
target = np.load(target_path, mmap_mode="r", allow_pickle=False)
|
||||
if target.ndim != 2 or target.shape[1] != 3:
|
||||
raise RuntimeError(
|
||||
f"Expected target shape [T, 3], found {target.shape}: "
|
||||
f"{target_path}"
|
||||
)
|
||||
|
||||
steps = int(target.shape[0])
|
||||
duration = steps / TARGET_FPS
|
||||
video_frames = math.ceil(steps * VIDEO_FPS / TARGET_FPS)
|
||||
specs.append(
|
||||
ClipSpec(
|
||||
split=split,
|
||||
uid=uid,
|
||||
steps=steps,
|
||||
duration=duration,
|
||||
video_frames=video_frames,
|
||||
source=source,
|
||||
target=target_path,
|
||||
)
|
||||
)
|
||||
return specs
|
||||
|
||||
|
||||
def write_manifest(specs: list[ClipSpec], path: Path) -> None:
|
||||
path.parent.mkdir(parents=True, exist_ok=True)
|
||||
temporary = path.with_suffix(path.suffix + ".partial")
|
||||
with temporary.open("w", newline="") as file:
|
||||
writer = csv.writer(file, delimiter="\t")
|
||||
writer.writerow(
|
||||
["split", "uid", "steps", "duration", "video_frames_2fps"]
|
||||
)
|
||||
for spec in specs:
|
||||
writer.writerow(
|
||||
[
|
||||
spec.split,
|
||||
spec.uid,
|
||||
spec.steps,
|
||||
f"{spec.duration:.6f}",
|
||||
spec.video_frames,
|
||||
]
|
||||
)
|
||||
temporary.replace(path)
|
||||
|
||||
|
||||
def select_specs(args: argparse.Namespace, specs: list[ClipSpec]) -> list[ClipSpec]:
|
||||
selected = specs
|
||||
if args.split != "all":
|
||||
selected = [spec for spec in selected if spec.split == args.split]
|
||||
if args.uid:
|
||||
requested = set(args.uid)
|
||||
known = {spec.uid for spec in specs}
|
||||
missing = sorted(requested - known)
|
||||
if missing:
|
||||
raise RuntimeError(
|
||||
"Requested UID(s) are not in EGO4DSHUFFLE: " + ", ".join(missing)
|
||||
)
|
||||
selected = [spec for spec in selected if spec.uid in requested]
|
||||
if args.limit is not None:
|
||||
selected = selected[: args.limit]
|
||||
if not selected:
|
||||
raise RuntimeError("No clips selected")
|
||||
return selected
|
||||
|
||||
|
||||
def count_video_frames(path: Path) -> int | None:
|
||||
if not path.is_file() or path.stat().st_size == 0:
|
||||
return None
|
||||
command = [
|
||||
"ffprobe",
|
||||
"-v",
|
||||
"error",
|
||||
"-select_streams",
|
||||
"v:0",
|
||||
"-count_frames",
|
||||
"-show_entries",
|
||||
"stream=nb_read_frames",
|
||||
"-of",
|
||||
"default=noprint_wrappers=1:nokey=1",
|
||||
str(path),
|
||||
]
|
||||
completed = subprocess.run(command, text=True, capture_output=True)
|
||||
if completed.returncode != 0:
|
||||
return None
|
||||
try:
|
||||
return int(completed.stdout.strip())
|
||||
except ValueError:
|
||||
return None
|
||||
|
||||
|
||||
def prepare_one(
|
||||
args: argparse.Namespace,
|
||||
spec: ClipSpec,
|
||||
index: int,
|
||||
total: int,
|
||||
) -> str:
|
||||
output = args.output_root / f"{spec.uid}_2fps.mp4"
|
||||
actual = count_video_frames(output)
|
||||
prefix = f"[{index}/{total}] [{spec.split}] {spec.uid}"
|
||||
|
||||
if not args.force and actual == spec.video_frames:
|
||||
print(f"{prefix}: reuse ({actual} frames)", flush=True)
|
||||
return "reused"
|
||||
|
||||
if actual is None:
|
||||
reason = "missing"
|
||||
else:
|
||||
reason = f"actual={actual}, expected={spec.video_frames}"
|
||||
print(
|
||||
f"{prefix}: prepare ({reason}; {spec.steps} states, "
|
||||
f"{spec.duration:.1f}s, {spec.video_frames} frames)",
|
||||
flush=True,
|
||||
)
|
||||
if args.dry_run:
|
||||
return "dry-run"
|
||||
|
||||
temporary = args.output_root / f".{spec.uid}.partial.mp4"
|
||||
duration = f"{spec.duration:.6f}"
|
||||
video_filter = (
|
||||
f"[0:v:0]trim=start=0:end={duration},setpts=PTS-STARTPTS,"
|
||||
f"fps=fps={VIDEO_FPS}:start_time=0:round=down,"
|
||||
"tpad=stop_mode=clone:stop=2,"
|
||||
f"trim=start_frame=0:end_frame={spec.video_frames},"
|
||||
f"setpts=N/({VIDEO_FPS}*TB),format=yuv420p[vout]"
|
||||
)
|
||||
audio_filter = (
|
||||
f"[0:a:0]apad,atrim=start=0:end={duration},"
|
||||
"asetpts=PTS-STARTPTS[aout]"
|
||||
)
|
||||
command = [
|
||||
"ffmpeg",
|
||||
"-hide_banner",
|
||||
"-loglevel",
|
||||
"warning",
|
||||
"-stats",
|
||||
"-y",
|
||||
"-i",
|
||||
str(spec.source),
|
||||
"-filter_complex",
|
||||
f"{video_filter};{audio_filter}",
|
||||
"-map",
|
||||
"[vout]",
|
||||
"-map",
|
||||
"[aout]",
|
||||
"-map_metadata",
|
||||
"-1",
|
||||
"-c:v",
|
||||
"libx264",
|
||||
"-preset",
|
||||
args.preset,
|
||||
"-crf",
|
||||
str(args.crf),
|
||||
"-c:a",
|
||||
"aac",
|
||||
"-ar",
|
||||
"16000",
|
||||
"-ac",
|
||||
"1",
|
||||
"-b:a",
|
||||
"64k",
|
||||
"-movflags",
|
||||
"+faststart",
|
||||
str(temporary),
|
||||
]
|
||||
completed = subprocess.run(command)
|
||||
if completed.returncode != 0:
|
||||
raise RuntimeError(
|
||||
f"ffmpeg failed for {spec.uid} with exit code {completed.returncode}; "
|
||||
f"partial output was left at {temporary}"
|
||||
)
|
||||
|
||||
actual = count_video_frames(temporary)
|
||||
if actual != spec.video_frames:
|
||||
raise RuntimeError(
|
||||
f"Prepared frame count is {actual}, expected {spec.video_frames}: "
|
||||
f"{temporary}"
|
||||
)
|
||||
temporary.replace(output)
|
||||
print(f"{prefix}: OK ({actual} frames) -> {output}", flush=True)
|
||||
return "prepared"
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
resolve_paths(args)
|
||||
for executable in ("ffmpeg", "ffprobe"):
|
||||
if shutil.which(executable) is None:
|
||||
raise RuntimeError(f"Required executable is not on PATH: {executable}")
|
||||
|
||||
specs = load_specs(args)
|
||||
write_manifest(specs, args.manifest)
|
||||
selected = select_specs(args, specs)
|
||||
args.output_root.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
total_steps = sum(spec.steps for spec in specs)
|
||||
print(
|
||||
f"Validated EGO4DSHUFFLE: {len(specs)} clips "
|
||||
f"({EXPECTED_TRAIN_CLIPS} train, {EXPECTED_TEST_CLIPS} test), "
|
||||
f"{total_steps} target states"
|
||||
)
|
||||
print(f"Manifest: {args.manifest}")
|
||||
print(f"Output: {args.output_root}")
|
||||
print(f"Selected: {len(selected)} clips; sequential ffmpeg execution")
|
||||
|
||||
counts = {"prepared": 0, "reused": 0, "dry-run": 0}
|
||||
for index, spec in enumerate(selected, start=1):
|
||||
result = prepare_one(args, spec, index, len(selected))
|
||||
counts[result] += 1
|
||||
|
||||
print(
|
||||
"Complete: "
|
||||
+ ", ".join(f"{key}={value}" for key, value in counts.items() if value)
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
try:
|
||||
main()
|
||||
except KeyboardInterrupt:
|
||||
print("\nInterrupted; rerun the same command to resume.", file=sys.stderr)
|
||||
raise SystemExit(130)
|
||||
except Exception as error:
|
||||
print(f"ERROR: {error}", file=sys.stderr)
|
||||
raise SystemExit(1)
|
||||
@@ -5,5 +5,7 @@ accelerate
|
||||
qwen-omni-utils[decord]
|
||||
librosa
|
||||
Pillow
|
||||
tqdm
|
||||
scikit-learn
|
||||
|
||||
# Optional on compatible NVIDIA GPUs: flash-attn
|
||||
|
||||
35
run_easycom_head_experiment.sh
Executable file
35
run_easycom_head_experiment.sh
Executable file
@@ -0,0 +1,35 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
# Frozen-Qwen head equivalent of EgoSpeak's run_LSTR_{train,inference}.sh.
|
||||
# Override any path or device by setting the corresponding environment value.
|
||||
ACTION="${1:-all}"
|
||||
if [[ "${ACTION}" != "train" && "${ACTION}" != "infer" && "${ACTION}" != "all" ]]; then
|
||||
echo "Usage: $0 [train|infer|all]" >&2
|
||||
exit 2
|
||||
fi
|
||||
PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/test_qwen_omni}"
|
||||
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/causal_prefix_2fps}"
|
||||
TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}"
|
||||
CHECKPOINT="${CHECKPOINT:-${PROJECT_ROOT}/checkpoints/causal_prefix_2fps_head.pth}"
|
||||
RESULT_ROOT="${RESULT_ROOT:-${PROJECT_ROOT}/results/causal_prefix_2fps_head}"
|
||||
DEVICE="${DEVICE:-cuda:0}"
|
||||
|
||||
cd "${PROJECT_ROOT}"
|
||||
|
||||
if [[ "${ACTION}" == "train" || "${ACTION}" == "all" ]]; then
|
||||
python train_turn_head.py \
|
||||
--feature-root "${FEATURE_ROOT}" \
|
||||
--target-root "${TARGET_ROOT}" \
|
||||
--output "${CHECKPOINT}" \
|
||||
--device "${DEVICE}"
|
||||
fi
|
||||
|
||||
if [[ "${ACTION}" == "infer" || "${ACTION}" == "all" ]]; then
|
||||
python infer_turn_head.py \
|
||||
--feature-root "${FEATURE_ROOT}" \
|
||||
--target-root "${TARGET_ROOT}" \
|
||||
--checkpoint "${CHECKPOINT}" \
|
||||
--output-dir "${RESULT_ROOT}" \
|
||||
--device "${DEVICE}"
|
||||
fi
|
||||
94
run_extract_5fps_2pergpu.sh
Executable file
94
run_extract_5fps_2pergpu.sh
Executable file
@@ -0,0 +1,94 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
MODEL="${MODEL:-/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B}"
|
||||
PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}"
|
||||
TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}"
|
||||
PREP_ROOT="${PREP_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_prepared}"
|
||||
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/causal_prefix_5fps}"
|
||||
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/causal_prefix_5fps}"
|
||||
|
||||
mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"
|
||||
cd "${PROJECT_ROOT}"
|
||||
|
||||
run_worker() {
|
||||
local gpu="$1"
|
||||
shift
|
||||
|
||||
local sid session metadata target output log
|
||||
for sid in "$@"; do
|
||||
session="Session_${sid}"
|
||||
metadata="${PREP_ROOT}/${session}/${session}.json"
|
||||
target="${TARGET_ROOT}/${session}.npy"
|
||||
output="${FEATURE_ROOT}/${session}.pt"
|
||||
log="${LOG_ROOT}/${session}.log"
|
||||
|
||||
if [[ ! -s "${metadata}" ]]; then
|
||||
echo "Missing metadata: ${metadata}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ ! -s "${target}" ]]; then
|
||||
echo "Missing target: ${target}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ -s "${output}" ]]; then
|
||||
echo "Skipping existing output: ${output}"
|
||||
continue
|
||||
fi
|
||||
|
||||
{
|
||||
echo "$(date): GPU ${gpu} starting ${session}"
|
||||
|
||||
CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
|
||||
python streaming_turn_model.py \
|
||||
--model-id "${MODEL}" \
|
||||
--session-metadata "${metadata}" \
|
||||
--target "${target}" \
|
||||
--use-audio-in-video \
|
||||
--mode causal-prefix \
|
||||
--memory-seconds 20 \
|
||||
--video-fps 5 \
|
||||
--strict-fixed-video-grid \
|
||||
--flash-attention-2 \
|
||||
--output "${output}"
|
||||
|
||||
echo "$(date): GPU ${gpu} completed ${session}"
|
||||
} >"${log}" 2>&1
|
||||
done
|
||||
}
|
||||
|
||||
pids=()
|
||||
|
||||
# Two independent extraction workers per GPU. Each GPU receives approximately
|
||||
# the same total number of EasyCom timesteps.
|
||||
run_worker 0 12 11 </dev/null &
|
||||
pids+=("$!")
|
||||
run_worker 0 7 </dev/null &
|
||||
pids+=("$!")
|
||||
|
||||
run_worker 1 6 10 </dev/null &
|
||||
pids+=("$!")
|
||||
run_worker 1 8 </dev/null &
|
||||
pids+=("$!")
|
||||
|
||||
run_worker 2 4 1 </dev/null &
|
||||
pids+=("$!")
|
||||
run_worker 2 9 </dev/null &
|
||||
pids+=("$!")
|
||||
|
||||
run_worker 3 5 2 </dev/null &
|
||||
pids+=("$!")
|
||||
run_worker 3 3 </dev/null &
|
||||
pids+=("$!")
|
||||
|
||||
status=0
|
||||
for pid in "${pids[@]}"; do
|
||||
wait "${pid}" || status=1
|
||||
done
|
||||
|
||||
if [[ "${status}" -ne 0 ]]; then
|
||||
echo "At least one extraction worker failed. Check ${LOG_ROOT}." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "ALL 12 CAUSAL 5 FPS EXTRACTIONS COMPLETED"
|
||||
134
run_extract_easycom_direct.sh
Normal file
134
run_extract_easycom_direct.sh
Normal file
@@ -0,0 +1,134 @@
|
||||
#!/usr/bin/env bash
|
||||
# Extract open-stream Thinker states for all 12 EasyCom sessions via
|
||||
# streaming_turn_model_directly.py (original clips, no prepared 2 FPS MP4).
|
||||
#
|
||||
# Recommended for full sessions: --memory-seconds 20 (bounded KV).
|
||||
# Do not default to --memory-seconds 0 on multi-minute sessions.
|
||||
set -euo pipefail
|
||||
|
||||
PROJECT_ROOT="${PROJECT_ROOT:-/mnt/58462F4C462F29E8/Code/icassp2027-experiment}"
|
||||
MODEL="${MODEL:-/mnt/58462F4C462F29E8/Models/Qwen2.5-Omni-3B}"
|
||||
VIDEO_ROOT="${VIDEO_ROOT:-/mnt/58462F4C462F29E8/Datasets/easycom/Main/Video_Compressed}"
|
||||
TARGET_ROOT="${TARGET_ROOT:-/mnt/58462F4C462F29E8/Code/EgoSpeak/dataset/EasyCom/target_perframe}"
|
||||
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct}"
|
||||
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct}"
|
||||
|
||||
MEMORY_SECONDS="${MEMORY_SECONDS:-20}"
|
||||
NATIVE_BLOCK_SECONDS="${NATIVE_BLOCK_SECONDS:-2}"
|
||||
VIDEO_FPS="${VIDEO_FPS:-2}"
|
||||
MODE="${MODE:-open-stream-replay}"
|
||||
DTYPE="${DTYPE:-bfloat16}"
|
||||
# Space-separated host GPU ids, e.g. "0" or "0 1 2 3". One session per GPU at a time.
|
||||
GPUS="${GPUS:-0}"
|
||||
# Optional: only run selected session ids, e.g. SESSIONS="1 2 3"
|
||||
SESSIONS="${SESSIONS:-1 2 3 4 5 6 7 8 9 10 11 12}"
|
||||
|
||||
# StreamConfig requires memory / native_block to be an integer (e.g. 20/2=10).
|
||||
# Values like 15 yield 7.5 and crash before loading the model.
|
||||
python - "${MEMORY_SECONDS}" "${NATIVE_BLOCK_SECONDS}" <<'PY'
|
||||
import sys
|
||||
memory, block = float(sys.argv[1]), float(sys.argv[2])
|
||||
ratio = memory / block
|
||||
if abs(ratio - round(ratio)) > 1e-6:
|
||||
raise SystemExit(
|
||||
f"MEMORY_SECONDS={memory} is invalid: memory/native_block={ratio} "
|
||||
f"(native_block={block}). Use a multiple of {block:g}, e.g. 16, 18, 20."
|
||||
)
|
||||
print(f"memory_seconds={memory:g} native_block={block:g} blocks={int(round(ratio))}")
|
||||
PY
|
||||
|
||||
mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"
|
||||
cd "${PROJECT_ROOT}"
|
||||
|
||||
run_one() {
|
||||
local gpu="$1"
|
||||
local sid="$2"
|
||||
local session="Session_${sid}"
|
||||
local session_dir="${VIDEO_ROOT}/${session}"
|
||||
local target="${TARGET_ROOT}/${session}.npy"
|
||||
local output="${FEATURE_ROOT}/${session}.pt"
|
||||
local log="${LOG_ROOT}/${session}.log"
|
||||
|
||||
if [[ ! -d "${session_dir}" ]]; then
|
||||
echo "Missing session dir: ${session_dir}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ ! -s "${target}" ]]; then
|
||||
echo "Missing target: ${target}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ -s "${output}" ]]; then
|
||||
echo "Skipping existing output: ${output}"
|
||||
return 0
|
||||
fi
|
||||
|
||||
{
|
||||
echo "$(date): GPU ${gpu} starting ${session}"
|
||||
echo "model=${MODEL}"
|
||||
echo "session_dir=${session_dir}"
|
||||
echo "target=${target}"
|
||||
echo "output=${output}"
|
||||
echo "memory_seconds=${MEMORY_SECONDS} video_fps=${VIDEO_FPS} mode=${MODE}"
|
||||
|
||||
if ! CUDA_VISIBLE_DEVICES="${gpu}" PYTHONUNBUFFERED=1 \
|
||||
python streaming_turn_model_directly.py \
|
||||
--model-id "${MODEL}" \
|
||||
--easycom-session-dir "${session_dir}" \
|
||||
--target "${target}" \
|
||||
--use-audio-in-video \
|
||||
--mode "${MODE}" \
|
||||
--memory-seconds "${MEMORY_SECONDS}" \
|
||||
--native-block-seconds "${NATIVE_BLOCK_SECONDS}" \
|
||||
--video-fps "${VIDEO_FPS}" \
|
||||
--dtype "${DTYPE}" \
|
||||
--flash-attention-2 \
|
||||
--device cuda:0 \
|
||||
--output "${output}"
|
||||
then
|
||||
echo "$(date): GPU ${gpu} FAILED ${session}"
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "$(date): GPU ${gpu} completed ${session}"
|
||||
} >"${log}" 2>&1
|
||||
}
|
||||
|
||||
# Round-robin sessions across GPUs.
|
||||
mapfile -t GPU_ARR < <(echo "${GPUS}" | tr ' ' '\n' | sed '/^$/d')
|
||||
mapfile -t SID_ARR < <(echo "${SESSIONS}" | tr ' ' '\n' | sed '/^$/d')
|
||||
|
||||
declare -A WORKERS=()
|
||||
for gpu in "${GPU_ARR[@]}"; do
|
||||
WORKERS["${gpu}"]=""
|
||||
done
|
||||
|
||||
i=0
|
||||
for sid in "${SID_ARR[@]}"; do
|
||||
gpu="${GPU_ARR[$((i % ${#GPU_ARR[@]}))]}"
|
||||
WORKERS["${gpu}"]+="${sid} "
|
||||
i=$((i + 1))
|
||||
done
|
||||
|
||||
pids=()
|
||||
status=0
|
||||
for gpu in "${GPU_ARR[@]}"; do
|
||||
sids="${WORKERS[${gpu}]}"
|
||||
[[ -z "${sids// }" ]] && continue
|
||||
(
|
||||
for sid in ${sids}; do
|
||||
run_one "${gpu}" "${sid}" || exit 1
|
||||
done
|
||||
) &
|
||||
pids+=("$!")
|
||||
done
|
||||
|
||||
for pid in "${pids[@]}"; do
|
||||
wait "${pid}" || status=1
|
||||
done
|
||||
|
||||
if [[ "${status}" -ne 0 ]]; then
|
||||
echo "At least one EasyCom direct extraction failed. Check ${LOG_ROOT}." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "ALL EasyCom direct extractions completed -> ${FEATURE_ROOT}"
|
||||
63
run_prepare_all_5fps.sh
Executable file
63
run_prepare_all_5fps.sh
Executable file
@@ -0,0 +1,63 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}"
|
||||
VIDEO_ROOT="${VIDEO_ROOT:-/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed}"
|
||||
TARGET_ROOT="${TARGET_ROOT:-${PROJECT_ROOT}/target_frame}"
|
||||
PREP_ROOT="${PREP_ROOT:-/data1/yfzx/wangx/Datasets/easycom/omni_prepared}"
|
||||
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/prepare_5fps}"
|
||||
|
||||
mkdir -p "${LOG_ROOT}"
|
||||
cd "${PROJECT_ROOT}"
|
||||
|
||||
prepare_worker() {
|
||||
local sid session session_dir target output_dir log
|
||||
for sid in "$@"; do
|
||||
session="Session_${sid}"
|
||||
session_dir="${VIDEO_ROOT}/${session}"
|
||||
target="${TARGET_ROOT}/${session}.npy"
|
||||
output_dir="${PREP_ROOT}/${session}"
|
||||
log="${LOG_ROOT}/${session}.log"
|
||||
|
||||
if [[ ! -d "${session_dir}" ]]; then
|
||||
echo "Missing source session: ${session_dir}" >&2
|
||||
return 1
|
||||
fi
|
||||
if [[ ! -s "${target}" ]]; then
|
||||
echo "Missing target: ${target}" >&2
|
||||
return 1
|
||||
fi
|
||||
|
||||
{
|
||||
echo "$(date): preparing ${session}"
|
||||
|
||||
python prepare_easycom_session.py \
|
||||
--session-dir "${session_dir}" \
|
||||
--output-dir "${output_dir}" \
|
||||
--target "${target}" \
|
||||
--fps 2 5 \
|
||||
--reuse-existing
|
||||
|
||||
echo "$(date): completed ${session}"
|
||||
} >"${log}" 2>&1
|
||||
done
|
||||
}
|
||||
|
||||
# FFmpeg/libx264 is already multithreaded. Two session workers generally make
|
||||
# better use of a many-core host without overwhelming storage or the CPU.
|
||||
prepare_worker 1 3 5 7 9 11 </dev/null &
|
||||
pid1=$!
|
||||
|
||||
prepare_worker 2 4 6 8 10 12 </dev/null &
|
||||
pid2=$!
|
||||
|
||||
status=0
|
||||
wait "${pid1}" || status=1
|
||||
wait "${pid2}" || status=1
|
||||
|
||||
if [[ "${status}" -ne 0 ]]; then
|
||||
echo "At least one preparation worker failed. Check ${LOG_ROOT}." >&2
|
||||
exit 1
|
||||
fi
|
||||
|
||||
echo "ALL 12 SESSIONS PREPARED AT 2 FPS AND 5 FPS"
|
||||
66
run_smoke_easycom_direct.sh
Normal file
66
run_smoke_easycom_direct.sh
Normal file
@@ -0,0 +1,66 @@
|
||||
#!/usr/bin/env bash
|
||||
# Smoke: one EasyCom session on one GPU via streaming_turn_model_directly.py
|
||||
#
|
||||
# Short smoke (default): --end only, NO --target
|
||||
# (passing both requires --end == full target duration)
|
||||
# Full session: END= (empty) and --target is used
|
||||
set -euo pipefail
|
||||
|
||||
# Remote defaults (override as needed)
|
||||
PROJECT_ROOT="${PROJECT_ROOT:-/data1/yfzx/wangx/Code/icassp2027-experiment}"
|
||||
MODEL="${MODEL:-/data1/yfzx/wangx/Models/Qwen2.5-Omni-3B}"
|
||||
VIDEO_ROOT="${VIDEO_ROOT:-/data1/yfzx/wangx/Datasets/easycom/extracted_easycom/Main/Video_Compressed}"
|
||||
TARGET_ROOT="${TARGET_ROOT:-/data1/yfzx/wangx/Datasets/easycom/target_frame}"
|
||||
FEATURE_ROOT="${FEATURE_ROOT:-${PROJECT_ROOT}/features/easycom_direct_smoke}"
|
||||
LOG_ROOT="${LOG_ROOT:-${PROJECT_ROOT}/logs/easycom_direct_smoke}"
|
||||
|
||||
SESSION_ID="${SESSION_ID:-1}"
|
||||
GPU="${GPU:-0}"
|
||||
# Short smoke: END=4. Full session smoke: END="" (empty).
|
||||
END="${END:-4}"
|
||||
MEMORY_SECONDS="${MEMORY_SECONDS:-20}"
|
||||
VIDEO_FPS="${VIDEO_FPS:-2}"
|
||||
|
||||
session="Session_${SESSION_ID}"
|
||||
session_dir="${VIDEO_ROOT}/${session}"
|
||||
target="${TARGET_ROOT}/${session}.npy"
|
||||
output="${FEATURE_ROOT}/${session}_gpu${GPU}.pt"
|
||||
log="${LOG_ROOT}/${session}_gpu${GPU}.log"
|
||||
|
||||
mkdir -p "${FEATURE_ROOT}" "${LOG_ROOT}"
|
||||
cd "${PROJECT_ROOT}"
|
||||
|
||||
[[ -d "${session_dir}" ]] || { echo "Missing: ${session_dir}" >&2; exit 1; }
|
||||
|
||||
extra_args=()
|
||||
if [[ -n "${END}" ]]; then
|
||||
# Bounded smoke: do not pass --target (length would conflict with --end).
|
||||
extra_args+=(--end "${END}")
|
||||
MEMORY_SECONDS=0
|
||||
else
|
||||
[[ -s "${target}" ]] || { echo "Missing: ${target}" >&2; exit 1; }
|
||||
extra_args+=(--target "${target}")
|
||||
fi
|
||||
|
||||
echo "GPU=${GPU} session=${session} end=${END:-full} memory=${MEMORY_SECONDS}"
|
||||
echo "log -> ${log}"
|
||||
|
||||
{
|
||||
echo "$(date): smoke start ${session} on GPU ${GPU}"
|
||||
CUDA_VISIBLE_DEVICES="${GPU}" PYTHONUNBUFFERED=1 \
|
||||
python streaming_turn_model_directly.py \
|
||||
--model-id "${MODEL}" \
|
||||
--easycom-session-dir "${session_dir}" \
|
||||
--use-audio-in-video \
|
||||
--mode open-stream-replay \
|
||||
--memory-seconds "${MEMORY_SECONDS}" \
|
||||
--video-fps "${VIDEO_FPS}" \
|
||||
--dtype bfloat16 \
|
||||
--flash-attention-2 \
|
||||
--device cuda:0 \
|
||||
--verify-every 5 \
|
||||
--verify-continuous-every 5 \
|
||||
"${extra_args[@]}" \
|
||||
--output "${output}"
|
||||
echo "$(date): smoke done -> ${output}"
|
||||
} 2>&1 | tee "${log}"
|
||||
@@ -4,24 +4,22 @@ This module implements two inference paths for Qwen2.5-Omni-3B:
|
||||
|
||||
1. ``causal-prefix``: recompute a causal AV window ending at every 200 ms.
|
||||
This is deliberately slow, but is the simplest correctness baseline.
|
||||
2. ``commit-replay``: keep the Thinker KV cache for completed 2-second
|
||||
multimodal blocks, and replay only the causal prefix of the current block.
|
||||
3. ``open-stream-replay``: cache the same completed AV blocks inside one
|
||||
2. ``open-stream-replay``: cache completed AV blocks inside one
|
||||
still-open user/media turn; closure and readout tokens remain temporary.
|
||||
|
||||
Neither path uses an RNN. A small head maps one Thinker state to the same
|
||||
``10 future horizons x 3 classes`` output used by EgoSpeak.
|
||||
|
||||
The commit/replay path touches version-specific Hugging Face internals. It is
|
||||
therefore pinned to Transformers 4.52.3 and includes an optional exact-replay
|
||||
check. Run that check before extracting a full dataset.
|
||||
The open-stream replay path touches version-specific Hugging Face internals.
|
||||
It is therefore pinned to Transformers 4.52.3 and includes an optional
|
||||
exact-replay check. Run that check before extracting a full dataset.
|
||||
|
||||
Example (no command is run merely by importing this file):
|
||||
|
||||
python streaming_turn_model.py \
|
||||
--video /path/to/session.mp4 \
|
||||
--use-audio-in-video \
|
||||
--mode commit-replay \
|
||||
--mode open-stream-replay \
|
||||
--memory-seconds 20 \
|
||||
--verify-every 25 \
|
||||
--output features/session.pt
|
||||
@@ -32,6 +30,7 @@ from __future__ import annotations
|
||||
import argparse
|
||||
import copy
|
||||
import json
|
||||
import logging
|
||||
import math
|
||||
import subprocess
|
||||
from dataclasses import asdict, dataclass
|
||||
@@ -39,16 +38,58 @@ from pathlib import Path
|
||||
from typing import Any, Literal
|
||||
|
||||
import torch
|
||||
from easycom_head import (
|
||||
TurnPredictionHead,
|
||||
egospeak_multihorizon_loss,
|
||||
miniroad_multihorizon_loss,
|
||||
multihorizon_cross_entropy,
|
||||
)
|
||||
from qwen_omni_utils import process_mm_info
|
||||
from torch import nn
|
||||
from tqdm.auto import tqdm
|
||||
from transformers import (
|
||||
Qwen2_5OmniProcessor,
|
||||
Qwen2_5OmniThinkerForConditionalGeneration,
|
||||
)
|
||||
from transformers.cache_utils import DynamicCache
|
||||
|
||||
Mode = Literal["causal-prefix", "commit-replay", "open-stream-replay"]
|
||||
Mode = Literal["causal-prefix", "open-stream-replay"]
|
||||
MODEL_ID = "Qwen/Qwen2.5-Omni-3B"
|
||||
_KNOWN_QWEN_WARNING_TEXT = (
|
||||
"System prompt modified, audio output may not work",
|
||||
"Unused or unrecognized kwargs: images.",
|
||||
"Unused or unrecognized kwargs: return_tensors, images.",
|
||||
"Unused or unrecognized kwargs: images, return_tensors.",
|
||||
)
|
||||
|
||||
|
||||
class _KnownQwenWarningFilter(logging.Filter):
|
||||
"""Hide Qwen's speech-output warning in Thinker-only extraction."""
|
||||
|
||||
def filter(self, record: logging.LogRecord) -> bool:
|
||||
message = record.getMessage()
|
||||
return not any(text in message for text in _KNOWN_QWEN_WARNING_TEXT)
|
||||
|
||||
|
||||
def _configure_qwen_logging() -> None:
|
||||
# Qwen2.5-Omni 4.52.3 emits this through Python's root logger whenever the
|
||||
# conversation does not use its speech-generation system prompt. This
|
||||
# extractor uses only the Thinker and never requests audio output.
|
||||
def add_filter(target: logging.Filterer) -> None:
|
||||
if not any(
|
||||
isinstance(item, _KnownQwenWarningFilter)
|
||||
for item in target.filters
|
||||
):
|
||||
target.addFilter(_KnownQwenWarningFilter())
|
||||
|
||||
# The system-prompt warning is logged directly on root. The unused-kwargs
|
||||
# warning is emitted by a Transformers child logger and handled by the
|
||||
# library's top-level StreamHandler, so cover both paths.
|
||||
for logger_name in (None, "transformers"):
|
||||
logger = logging.getLogger(logger_name)
|
||||
add_filter(logger)
|
||||
for handler in logger.handlers:
|
||||
add_filter(handler)
|
||||
|
||||
|
||||
def _choose_dtype(
|
||||
@@ -132,6 +173,7 @@ class OpenMediaPieces:
|
||||
payload_embeddings: torch.Tensor
|
||||
payload_position_ids: torch.LongTensor
|
||||
closure_ids: torch.LongTensor
|
||||
closure_position_ids: torch.LongTensor
|
||||
|
||||
|
||||
class FixedGridAV:
|
||||
@@ -148,10 +190,10 @@ class FixedGridAV:
|
||||
def __init__(self, path: str | Path, expected_fps: float) -> None:
|
||||
try:
|
||||
import decord
|
||||
import librosa
|
||||
import numpy as np
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"Fixed-grid AV extraction requires decord and librosa"
|
||||
"Fixed-grid AV extraction requires decord and numpy"
|
||||
) from exc
|
||||
|
||||
self.path = Path(path).resolve()
|
||||
@@ -162,11 +204,46 @@ class FixedGridAV:
|
||||
f"Prepared video is {self.fps:g} FPS, expected {expected_fps:g}"
|
||||
)
|
||||
self.num_frames = len(self.reader)
|
||||
self.audio, _ = librosa.load(
|
||||
str(self.path), sr=self.sample_rate, mono=True
|
||||
# SoundFile cannot decode the AAC track in MP4, causing librosa to
|
||||
# fall back to its deprecated and much slower audioread path. Decode
|
||||
# directly to the exact float32/16-kHz/mono representation consumed by
|
||||
# Qwen. Keeping the complete track in memory also makes repeated
|
||||
# 200-ms slicing deterministic and inexpensive.
|
||||
completed = subprocess.run(
|
||||
[
|
||||
"ffmpeg",
|
||||
"-nostdin",
|
||||
"-v",
|
||||
"error",
|
||||
"-i",
|
||||
str(self.path),
|
||||
"-map",
|
||||
"0:a:0",
|
||||
"-vn",
|
||||
"-ac",
|
||||
"1",
|
||||
"-ar",
|
||||
str(self.sample_rate),
|
||||
"-c:a",
|
||||
"pcm_f32le",
|
||||
"-f",
|
||||
"f32le",
|
||||
"pipe:1",
|
||||
],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
)
|
||||
if completed.returncode != 0:
|
||||
message = completed.stderr.decode("utf-8", errors="replace").strip()
|
||||
raise RuntimeError(
|
||||
f"Failed to decode embedded audio from {self.path}: {message}"
|
||||
)
|
||||
self.audio = np.frombuffer(completed.stdout, dtype="<f4").copy()
|
||||
if not len(self.audio):
|
||||
raise RuntimeError(f"Decoded audio is empty: {self.path}")
|
||||
|
||||
def frames(self, start: float, end: float) -> list[Any]:
|
||||
<<<<<<< HEAD
|
||||
"""Return fixed-grid frames whose timestamps lie in ``[start, end)``.
|
||||
|
||||
Half-open ownership is essential at native block boundaries. For
|
||||
@@ -174,18 +251,30 @@ class FixedGridAV:
|
||||
not also be committed in ``[0.0, 2.0)``. Audio slicing below follows
|
||||
the same convention.
|
||||
"""
|
||||
=======
|
||||
"""Return fixed-grid frames whose timestamps lie in [start, end)."""
|
||||
>>>>>>> 8145be410866dd0675038d92347aea34d24b84d5
|
||||
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
|
||||
first = max(0, int(math.ceil(start * self.fps - 1e-7)))
|
||||
<<<<<<< HEAD
|
||||
last = min(
|
||||
self.num_frames - 1,
|
||||
int(math.ceil(end * self.fps - 1e-7)) - 1,
|
||||
=======
|
||||
last_exclusive = min(
|
||||
self.num_frames,
|
||||
int(math.ceil(end * self.fps - 1e-7)),
|
||||
>>>>>>> 8145be410866dd0675038d92347aea34d24b84d5
|
||||
)
|
||||
last = last_exclusive - 1
|
||||
if first > last:
|
||||
# This can occur only when a caller starts between global visual
|
||||
# ticks. Reuse the most recent past frame, never ceil into future.
|
||||
# A short interval can fall entirely between global visual ticks.
|
||||
# Reuse the latest frame strictly before ``end``, never a future
|
||||
# frame. Native 2-second blocks start on-grid and do not take this
|
||||
# fallback path.
|
||||
first = last = max(0, min(self.num_frames - 1, last))
|
||||
indices = list(range(first, last + 1))
|
||||
arrays = self.reader.get_batch(indices).asnumpy()
|
||||
@@ -204,76 +293,6 @@ class FixedGridAV:
|
||||
return self.audio[first:last]
|
||||
|
||||
|
||||
class TurnPredictionHead(nn.Module):
|
||||
"""Map one Omni state to EgoSpeak-compatible [10, 3] logits."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
hidden_size: int = 2048,
|
||||
num_horizons: int = 10,
|
||||
num_classes: int = 3,
|
||||
dropout: float = 0.1,
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.num_horizons = num_horizons
|
||||
self.num_classes = num_classes
|
||||
self.net = nn.Sequential(
|
||||
nn.LayerNorm(hidden_size),
|
||||
nn.Linear(hidden_size, hidden_size // 2),
|
||||
nn.GELU(),
|
||||
nn.Dropout(dropout),
|
||||
nn.Linear(hidden_size // 2, num_horizons * num_classes),
|
||||
)
|
||||
|
||||
def forward(self, states: torch.Tensor) -> torch.Tensor:
|
||||
logits = self.net(states)
|
||||
return logits.view(*states.shape[:-1], self.num_horizons, self.num_classes)
|
||||
|
||||
|
||||
def egospeak_multihorizon_loss(
|
||||
logits: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
*,
|
||||
reduction: Literal["mean", "sum", "none"] = "mean",
|
||||
) -> torch.Tensor:
|
||||
"""Reproduce EgoSpeak's anticipation loss for multi-positive targets.
|
||||
|
||||
``logits`` and ``targets`` both have shape ``[..., 10, 3]``. EasyCom can
|
||||
label camera-wearer and other speech simultaneously as ``[0, 1, 1]``.
|
||||
MiniROAD L2-normalizes that row and applies a log-softmax loss.
|
||||
"""
|
||||
|
||||
if logits.shape != targets.shape:
|
||||
raise ValueError(
|
||||
f"logits/targets shape mismatch: {tuple(logits.shape)} vs "
|
||||
f"{tuple(targets.shape)}"
|
||||
)
|
||||
if logits.shape[-1] != 3:
|
||||
raise ValueError(f"Expected three classes, got {logits.shape[-1]}")
|
||||
targets = nn.functional.normalize(
|
||||
targets.to(device=logits.device, dtype=logits.dtype),
|
||||
p=2,
|
||||
dim=-1,
|
||||
)
|
||||
loss = -(targets * nn.functional.log_softmax(logits, dim=-1)).sum(dim=-1)
|
||||
if reduction == "mean":
|
||||
return loss.mean()
|
||||
if reduction == "sum":
|
||||
return loss.sum()
|
||||
if reduction == "none":
|
||||
return loss
|
||||
raise ValueError(f"Unknown reduction: {reduction}")
|
||||
|
||||
|
||||
def multihorizon_cross_entropy(
|
||||
logits: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
) -> torch.Tensor:
|
||||
"""Backward-compatible name for :func:`egospeak_multihorizon_loss`."""
|
||||
|
||||
return egospeak_multihorizon_loss(logits, targets)
|
||||
|
||||
|
||||
def _require_integer_ratio(numerator: float, denominator: float, name: str) -> int:
|
||||
ratio = numerator / denominator
|
||||
rounded = round(ratio)
|
||||
@@ -443,16 +462,23 @@ def _prepare_conversation(
|
||||
# Qwen2.5-Omni uses this value to compute video_second_per_grid for
|
||||
# TMRoPE. This is essential for the 5 FPS ablation (default is 2 FPS).
|
||||
processor_kwargs["fps"] = video_fps
|
||||
batch = processor(
|
||||
text=[text],
|
||||
audio=audios,
|
||||
images=images,
|
||||
videos=videos,
|
||||
padding=True,
|
||||
return_tensors="pt",
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
processor_inputs: dict[str, Any] = {
|
||||
"text": [text],
|
||||
"audio": audios,
|
||||
"videos": videos,
|
||||
"padding": True,
|
||||
"use_audio_in_video": use_audio_in_video,
|
||||
**processor_kwargs,
|
||||
)
|
||||
}
|
||||
# Passing images=None and return_tensors through Qwen's 4.52.3 processor
|
||||
# produces a warning on every streaming step. Omit the absent modality and
|
||||
# perform the same BatchFeature tensor conversion explicitly afterward.
|
||||
if images is not None and (
|
||||
not isinstance(images, (list, tuple)) or len(images) > 0
|
||||
):
|
||||
processor_inputs["images"] = images
|
||||
batch = processor(**processor_inputs)
|
||||
batch = batch.convert_to_tensors("pt")
|
||||
# BatchFeature.to(dtype) changes floating tensors only; IDs stay integral.
|
||||
batch = batch.to(model.device).to(model.dtype)
|
||||
return dict(batch)
|
||||
@@ -668,18 +694,10 @@ def _open_media_pieces(
|
||||
payload_embeddings=embeddings[:, first : last + 1, :],
|
||||
payload_position_ids=positions[:, :, first : last + 1],
|
||||
closure_ids=closure_ids,
|
||||
closure_position_ids=positions[:, :, last + 1 :],
|
||||
)
|
||||
|
||||
|
||||
def _positions_at_offset(
|
||||
position_ids: torch.LongTensor,
|
||||
offset: int,
|
||||
) -> torch.LongTensor:
|
||||
"""Preserve local 3D geometry while moving its minimum to ``offset``."""
|
||||
|
||||
return position_ids - position_ids.min() + offset
|
||||
|
||||
|
||||
def _tmrope_positions_per_chunk(
|
||||
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
||||
) -> int:
|
||||
@@ -829,13 +847,22 @@ def _read_open_state(
|
||||
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
||||
open_media_cache: Any,
|
||||
*,
|
||||
next_rope_offset: int,
|
||||
closure_ids: torch.LongTensor,
|
||||
closure_position_ids: torch.LongTensor,
|
||||
query_ids: torch.LongTensor,
|
||||
) -> torch.Tensor:
|
||||
"""Temporarily close an open media turn, read q_t, and discard the branch."""
|
||||
|
||||
branch_cache = clone_cache(open_media_cache)
|
||||
if closure_position_ids.shape != (
|
||||
3,
|
||||
closure_ids.shape[0],
|
||||
closure_ids.shape[1],
|
||||
):
|
||||
raise ValueError(
|
||||
"closure position/ID mismatch: "
|
||||
f"{tuple(closure_position_ids.shape)} vs {tuple(closure_ids.shape)}"
|
||||
)
|
||||
readout_ids = torch.cat(
|
||||
[closure_ids.to(query_ids.device), query_ids],
|
||||
dim=1,
|
||||
@@ -843,13 +870,24 @@ def _read_open_state(
|
||||
past_length = _cache_length(branch_cache)
|
||||
readout_length = readout_ids.shape[1]
|
||||
embeddings = model.get_input_embeddings()(readout_ids)
|
||||
scalar_positions = torch.arange(
|
||||
next_rope_offset,
|
||||
next_rope_offset + readout_length,
|
||||
|
||||
# Qwen's AV TMRoPE gives the paired video/audio closing tokens the same
|
||||
# position. Preserve those model-derived positions instead of replacing
|
||||
# the whole closure with a scalar arange. Only the artificial assistant
|
||||
# query follows the closure's maximum position.
|
||||
closure_position_ids = closure_position_ids.to(query_ids.device)
|
||||
query_start = int(closure_position_ids.max().item()) + 1
|
||||
query_scalar_positions = torch.arange(
|
||||
query_start,
|
||||
query_start + query_ids.shape[1],
|
||||
device=readout_ids.device,
|
||||
)
|
||||
position_ids = scalar_positions.view(1, 1, -1).expand(
|
||||
3, readout_ids.shape[0], -1
|
||||
query_position_ids = query_scalar_positions.view(1, 1, -1).expand(
|
||||
3, query_ids.shape[0], -1
|
||||
)
|
||||
position_ids = torch.cat(
|
||||
[closure_position_ids, query_position_ids],
|
||||
dim=2,
|
||||
)
|
||||
cache_position = torch.arange(
|
||||
past_length,
|
||||
@@ -882,32 +920,41 @@ def _append_open_pieces(
|
||||
past_key_values: Any | None,
|
||||
rope_offset: int,
|
||||
include_prefix: bool,
|
||||
) -> tuple[Any, int]:
|
||||
"""Append one AV payload, adding the chat/media prefix exactly once."""
|
||||
) -> tuple[Any, int, torch.LongTensor]:
|
||||
"""Append one AV payload and return its shifted closure positions."""
|
||||
|
||||
if include_prefix:
|
||||
if past_key_values is not None:
|
||||
raise ValueError("An open-media prefix can only start an empty cache")
|
||||
payload_base = int(pieces.payload_position_ids.min().item())
|
||||
if rope_offset != payload_base:
|
||||
raise ValueError(
|
||||
"The first open payload must retain its local TMRoPE base: "
|
||||
f"got {rope_offset}, expected {payload_base}"
|
||||
)
|
||||
embeddings = torch.cat(
|
||||
[pieces.prefix_embeddings, pieces.payload_embeddings], dim=1
|
||||
)
|
||||
position_ids = torch.cat(
|
||||
[pieces.prefix_position_ids, pieces.payload_position_ids], dim=2
|
||||
)
|
||||
closure_position_ids = pieces.closure_position_ids
|
||||
else:
|
||||
if past_key_values is None:
|
||||
raise ValueError("Payload-only append requires an existing open cache")
|
||||
embeddings = pieces.payload_embeddings
|
||||
position_ids = _positions_at_offset(
|
||||
pieces.payload_position_ids,
|
||||
rope_offset,
|
||||
position_shift = rope_offset - int(
|
||||
pieces.payload_position_ids.min().item()
|
||||
)
|
||||
return _append_embeddings(
|
||||
position_ids = pieces.payload_position_ids + position_shift
|
||||
closure_position_ids = pieces.closure_position_ids + position_shift
|
||||
next_cache, next_rope_offset = _append_embeddings(
|
||||
model,
|
||||
embeddings,
|
||||
position_ids,
|
||||
past_key_values=past_key_values,
|
||||
)
|
||||
return next_cache, next_rope_offset, closure_position_ids
|
||||
|
||||
|
||||
class CausalPrefixExtractor:
|
||||
@@ -988,18 +1035,34 @@ class CausalPrefixExtractor:
|
||||
start: float,
|
||||
end: float,
|
||||
use_audio_in_video: bool,
|
||||
show_progress: bool = True,
|
||||
) -> ExtractionResult:
|
||||
times = _time_grid(start, end, self.config.step_seconds)
|
||||
states = []
|
||||
for current_time in times:
|
||||
context_start = max(start, current_time - self.config.memory_seconds)
|
||||
state = self.state_for_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=[(context_start, current_time)],
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
states.append(state.cpu())
|
||||
with tqdm(
|
||||
total=len(times),
|
||||
desc="causal-prefix",
|
||||
unit="state",
|
||||
dynamic_ncols=True,
|
||||
mininterval=0.0,
|
||||
miniters=1,
|
||||
disable=not show_progress,
|
||||
) as progress:
|
||||
for current_time in times:
|
||||
context_start = max(
|
||||
start, current_time - self.config.memory_seconds
|
||||
)
|
||||
state = self.state_for_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=[(context_start, current_time)],
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
states.append(state.cpu())
|
||||
progress.set_postfix_str(
|
||||
f"media={current_time:.1f}s", refresh=False
|
||||
)
|
||||
progress.update(1)
|
||||
return ExtractionResult(
|
||||
times=torch.tensor(times, dtype=torch.float64),
|
||||
states=torch.cat(states, dim=0),
|
||||
@@ -1007,202 +1070,6 @@ class CausalPrefixExtractor:
|
||||
)
|
||||
|
||||
|
||||
class CommitReplayExtractor:
|
||||
"""Cache complete 2 s blocks and replay only the current causal prefix."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
model: Qwen2_5OmniThinkerForConditionalGeneration,
|
||||
processor: Qwen2_5OmniProcessor,
|
||||
config: StreamConfig,
|
||||
fixed_grid_av: FixedGridAV | None = None,
|
||||
) -> None:
|
||||
config.validate()
|
||||
self.model = model
|
||||
self.processor = processor
|
||||
self.config = config
|
||||
self.fixed_grid_av = fixed_grid_av
|
||||
self.query_ids = _readout_query_ids(processor, model.device)
|
||||
self.reference = CausalPrefixExtractor(
|
||||
model, processor, config, fixed_grid_av=fixed_grid_av
|
||||
)
|
||||
|
||||
def _prepare_one(
|
||||
self,
|
||||
*,
|
||||
video_path: str | Path | None,
|
||||
audio_path: str | Path | None,
|
||||
interval: tuple[float, float],
|
||||
use_audio_in_video: bool,
|
||||
turn_index: int,
|
||||
) -> dict[str, torch.Tensor]:
|
||||
conversation, manual_audios = build_chunk_conversation(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=[interval],
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
config=self.config,
|
||||
fixed_grid_av=self.fixed_grid_av,
|
||||
)
|
||||
text_override = _media_turn_template_suffix(
|
||||
self.processor,
|
||||
conversation[0],
|
||||
turn_index=turn_index,
|
||||
)
|
||||
return _prepare_conversation(
|
||||
self.processor,
|
||||
conversation,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
model=self.model,
|
||||
manual_audios=manual_audios,
|
||||
video_fps=self.config.video_fps,
|
||||
text_override=text_override,
|
||||
)
|
||||
|
||||
def _extract_segment(
|
||||
self,
|
||||
*,
|
||||
video_path: str | Path | None,
|
||||
audio_path: str | Path | None,
|
||||
segment_start: float,
|
||||
segment_end: float,
|
||||
use_audio_in_video: bool,
|
||||
verify_every: int,
|
||||
verify_continuous_every: int,
|
||||
global_step_offset: int,
|
||||
) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]:
|
||||
times = _time_grid(segment_start, segment_end, self.config.step_seconds)
|
||||
committed_cache = None
|
||||
committed_rope_offset = 0
|
||||
block_start = segment_start
|
||||
states: list[torch.Tensor] = []
|
||||
checks: list[dict[str, Any]] = []
|
||||
|
||||
for local_step, current_time in enumerate(times, start=1):
|
||||
while current_time > block_start + self.config.native_block_seconds + 1e-6:
|
||||
block_start += self.config.native_block_seconds
|
||||
|
||||
# DynamicCache is mutated by forward. Every unfinished prefix must
|
||||
# branch from the identical committed history.
|
||||
branch_cache = clone_cache(committed_cache)
|
||||
inputs = self._prepare_one(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
interval=(block_start, current_time),
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
turn_index=round(
|
||||
(block_start - segment_start)
|
||||
/ self.config.native_block_seconds
|
||||
),
|
||||
)
|
||||
media_cache, next_rope_offset = _append_media_turn(
|
||||
self.model,
|
||||
inputs,
|
||||
past_key_values=branch_cache,
|
||||
rope_offset=committed_rope_offset,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
state = _read_state(
|
||||
self.model,
|
||||
media_cache,
|
||||
next_rope_offset=next_rope_offset,
|
||||
query_ids=self.query_ids,
|
||||
)
|
||||
states.append(state.cpu())
|
||||
|
||||
global_step = global_step_offset + local_step
|
||||
if verify_every > 0 and global_step % verify_every == 0:
|
||||
intervals = _block_intervals(
|
||||
segment_start,
|
||||
current_time,
|
||||
self.config.native_block_seconds,
|
||||
)
|
||||
exact = self.reference.state_for_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=intervals,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
checks.append(
|
||||
_compare_states(current_time, state, exact, kind="cache_replay")
|
||||
)
|
||||
if (
|
||||
verify_continuous_every > 0
|
||||
and global_step % verify_continuous_every == 0
|
||||
):
|
||||
continuous = self.reference.state_for_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=[(segment_start, current_time)],
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
)
|
||||
checks.append(
|
||||
_compare_states(
|
||||
current_time,
|
||||
state,
|
||||
continuous,
|
||||
kind="continuous_window_semantics",
|
||||
)
|
||||
)
|
||||
|
||||
block_end = min(
|
||||
block_start + self.config.native_block_seconds,
|
||||
segment_end,
|
||||
)
|
||||
if math.isclose(current_time, block_end, rel_tol=0.0, abs_tol=1e-6):
|
||||
# media_cache excludes the readout query, so it is safe to keep.
|
||||
committed_cache = media_cache
|
||||
committed_rope_offset = next_rope_offset
|
||||
block_start = block_end
|
||||
|
||||
return times, states, checks
|
||||
|
||||
def extract(
|
||||
self,
|
||||
*,
|
||||
video_path: str | Path | None,
|
||||
audio_path: str | Path | None,
|
||||
start: float,
|
||||
end: float,
|
||||
use_audio_in_video: bool,
|
||||
verify_every: int = 0,
|
||||
verify_continuous_every: int = 0,
|
||||
) -> ExtractionResult:
|
||||
all_times: list[float] = []
|
||||
all_states: list[torch.Tensor] = []
|
||||
all_checks: list[dict[str, Any]] = []
|
||||
segment_start = start
|
||||
global_steps = 0
|
||||
|
||||
# A bounded cache is implemented as non-overlapping memory segments.
|
||||
# Predictions are still emitted for the entire session, but history is
|
||||
# reset at each boundary. Rolling KV eviction is intentionally not
|
||||
# attempted because it also requires rebasing multimodal RoPE.
|
||||
while segment_start < end - 1e-7:
|
||||
segment_end = min(segment_start + self.config.memory_seconds, end)
|
||||
times, states, checks = self._extract_segment(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
segment_start=segment_start,
|
||||
segment_end=segment_end,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
verify_every=verify_every,
|
||||
verify_continuous_every=verify_continuous_every,
|
||||
global_step_offset=global_steps,
|
||||
)
|
||||
all_times.extend(times)
|
||||
all_states.extend(states)
|
||||
all_checks.extend(checks)
|
||||
global_steps += len(times)
|
||||
segment_start = segment_end
|
||||
|
||||
return ExtractionResult(
|
||||
times=torch.tensor(all_times, dtype=torch.float64),
|
||||
states=torch.cat(all_states, dim=0),
|
||||
verification=all_checks,
|
||||
)
|
||||
|
||||
|
||||
class OpenStreamReplayExtractor:
|
||||
"""Cache completed AV payloads inside one still-open multimodal turn."""
|
||||
|
||||
@@ -1298,7 +1165,13 @@ class OpenStreamReplayExtractor:
|
||||
) -> torch.Tensor:
|
||||
"""Replay the same open stream from an empty cache for verification."""
|
||||
|
||||
media_cache, next_offset, _, closure_ids = self._cache_from_intervals(
|
||||
(
|
||||
media_cache,
|
||||
_,
|
||||
_,
|
||||
closure_ids,
|
||||
closure_position_ids,
|
||||
) = self._cache_from_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
intervals=intervals,
|
||||
@@ -1307,8 +1180,8 @@ class OpenStreamReplayExtractor:
|
||||
return _read_open_state(
|
||||
self.model,
|
||||
media_cache,
|
||||
next_rope_offset=next_offset,
|
||||
closure_ids=closure_ids,
|
||||
closure_position_ids=closure_position_ids,
|
||||
query_ids=self.query_ids,
|
||||
)
|
||||
|
||||
@@ -1319,7 +1192,7 @@ class OpenStreamReplayExtractor:
|
||||
audio_path: str | Path | None,
|
||||
intervals: list[tuple[float, float]],
|
||||
use_audio_in_video: bool,
|
||||
) -> tuple[Any, int, int, torch.LongTensor]:
|
||||
) -> tuple[Any, int, int, torch.LongTensor, torch.LongTensor]:
|
||||
"""Build an open payload cache and expose its local TMRoPE base."""
|
||||
|
||||
if not intervals:
|
||||
@@ -1327,6 +1200,7 @@ class OpenStreamReplayExtractor:
|
||||
media_cache = None
|
||||
next_offset = 0
|
||||
closure_ids = None
|
||||
closure_position_ids = None
|
||||
payload_base = None
|
||||
positions_per_chunk = _tmrope_positions_per_chunk(self.model)
|
||||
for index, interval in enumerate(intervals):
|
||||
@@ -1339,7 +1213,11 @@ class OpenStreamReplayExtractor:
|
||||
if payload_base is None:
|
||||
payload_base = int(pieces.payload_position_ids.min().item())
|
||||
payload_offset = payload_base + index * positions_per_chunk
|
||||
media_cache, next_offset = _append_open_pieces(
|
||||
(
|
||||
media_cache,
|
||||
next_offset,
|
||||
closure_position_ids,
|
||||
) = _append_open_pieces(
|
||||
self.model,
|
||||
pieces,
|
||||
past_key_values=media_cache,
|
||||
@@ -1349,9 +1227,17 @@ class OpenStreamReplayExtractor:
|
||||
closure_ids = pieces.closure_ids
|
||||
if closure_ids is None:
|
||||
raise RuntimeError("Open-stream replay did not produce closure tokens")
|
||||
if closure_position_ids is None:
|
||||
raise RuntimeError("Open-stream replay did not produce closure positions")
|
||||
if payload_base is None:
|
||||
raise RuntimeError("Open-stream replay did not establish a payload base")
|
||||
return media_cache, next_offset, payload_base, closure_ids
|
||||
return (
|
||||
media_cache,
|
||||
next_offset,
|
||||
payload_base,
|
||||
closure_ids,
|
||||
closure_position_ids,
|
||||
)
|
||||
|
||||
def _extract_segment(
|
||||
self,
|
||||
@@ -1364,6 +1250,7 @@ class OpenStreamReplayExtractor:
|
||||
verify_every: int,
|
||||
verify_continuous_every: int,
|
||||
global_step_offset: int,
|
||||
progress: Any,
|
||||
) -> tuple[list[float], list[torch.Tensor], list[dict[str, Any]]]:
|
||||
times = _time_grid(segment_start, segment_end, self.config.step_seconds)
|
||||
committed_cache = None
|
||||
@@ -1400,6 +1287,7 @@ class OpenStreamReplayExtractor:
|
||||
_,
|
||||
payload_base,
|
||||
_,
|
||||
_,
|
||||
) = self._cache_from_intervals(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
@@ -1421,7 +1309,11 @@ class OpenStreamReplayExtractor:
|
||||
)
|
||||
payload_offset = payload_base + block_index * positions_per_chunk
|
||||
branch_cache = clone_cache(committed_cache)
|
||||
media_cache, next_offset = _append_open_pieces(
|
||||
(
|
||||
media_cache,
|
||||
_,
|
||||
closure_position_ids,
|
||||
) = _append_open_pieces(
|
||||
self.model,
|
||||
pieces,
|
||||
past_key_values=branch_cache,
|
||||
@@ -1431,8 +1323,8 @@ class OpenStreamReplayExtractor:
|
||||
state = _read_open_state(
|
||||
self.model,
|
||||
media_cache,
|
||||
next_rope_offset=next_offset,
|
||||
closure_ids=pieces.closure_ids,
|
||||
closure_position_ids=closure_position_ids,
|
||||
query_ids=self.query_ids,
|
||||
)
|
||||
states.append(state.cpu())
|
||||
@@ -1481,6 +1373,11 @@ class OpenStreamReplayExtractor:
|
||||
committed_cache = media_cache
|
||||
block_start = block_end
|
||||
|
||||
progress.set_postfix_str(
|
||||
f"media={current_time:.1f}s", refresh=False
|
||||
)
|
||||
progress.update(1)
|
||||
|
||||
return times, states, checks
|
||||
|
||||
def extract(
|
||||
@@ -1493,17 +1390,29 @@ class OpenStreamReplayExtractor:
|
||||
use_audio_in_video: bool,
|
||||
verify_every: int = 0,
|
||||
verify_continuous_every: int = 0,
|
||||
show_progress: bool = True,
|
||||
) -> ExtractionResult:
|
||||
all_times, all_states, all_checks = self._extract_segment(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
segment_start=start,
|
||||
segment_end=end,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
verify_every=verify_every,
|
||||
verify_continuous_every=verify_continuous_every,
|
||||
global_step_offset=0,
|
||||
)
|
||||
total_steps = len(_time_grid(start, end, self.config.step_seconds))
|
||||
with tqdm(
|
||||
total=total_steps,
|
||||
desc="open-stream-replay",
|
||||
unit="state",
|
||||
dynamic_ncols=True,
|
||||
mininterval=0.0,
|
||||
miniters=1,
|
||||
disable=not show_progress,
|
||||
) as progress:
|
||||
all_times, all_states, all_checks = self._extract_segment(
|
||||
video_path=video_path,
|
||||
audio_path=audio_path,
|
||||
segment_start=start,
|
||||
segment_end=end,
|
||||
use_audio_in_video=use_audio_in_video,
|
||||
verify_every=verify_every,
|
||||
verify_continuous_every=verify_continuous_every,
|
||||
global_step_offset=0,
|
||||
progress=progress,
|
||||
)
|
||||
return ExtractionResult(
|
||||
times=torch.tensor(all_times, dtype=torch.float64),
|
||||
states=torch.cat(all_states, dim=0),
|
||||
@@ -1512,7 +1421,7 @@ class OpenStreamReplayExtractor:
|
||||
|
||||
|
||||
def _block_intervals(start: float, end: float, block_seconds: float) -> list[tuple[float, float]]:
|
||||
"""Exact media-turn sequence represented by a commit/replay cache."""
|
||||
"""Exact AV payload intervals represented by an open-stream cache."""
|
||||
|
||||
intervals = []
|
||||
cursor = start
|
||||
@@ -1664,7 +1573,7 @@ def parse_args() -> argparse.Namespace:
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=("causal-prefix", "commit-replay", "open-stream-replay"),
|
||||
choices=("causal-prefix", "open-stream-replay"),
|
||||
default="open-stream-replay",
|
||||
)
|
||||
parser.add_argument("--model-id", default=MODEL_ID)
|
||||
@@ -1688,6 +1597,11 @@ def parse_args() -> argparse.Namespace:
|
||||
default=0,
|
||||
help="Also compare block replay with one continuous causal media window.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--no-progress",
|
||||
action="store_true",
|
||||
help="Disable the per-state tqdm progress bar.",
|
||||
)
|
||||
parser.add_argument("--head-checkpoint", type=Path)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
@@ -1701,6 +1615,7 @@ def parse_args() -> argparse.Namespace:
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
_configure_qwen_logging()
|
||||
session_metadata = _apply_session_metadata(args)
|
||||
if args.video is None and args.audio is None:
|
||||
raise SystemExit("Pass --video, --audio, or both")
|
||||
@@ -1710,7 +1625,6 @@ def main() -> None:
|
||||
raise SystemExit("--verify-every must be >= 0")
|
||||
if args.verify_continuous_every < 0:
|
||||
raise SystemExit("--verify-continuous-every must be >= 0")
|
||||
|
||||
if args.num_steps is not None and args.num_steps <= 0:
|
||||
raise SystemExit("--num-steps must be positive")
|
||||
if args.target is not None:
|
||||
@@ -1785,19 +1699,7 @@ def main() -> None:
|
||||
start=args.start,
|
||||
end=end,
|
||||
use_audio_in_video=args.use_audio_in_video,
|
||||
)
|
||||
elif args.mode == "commit-replay":
|
||||
extractor = CommitReplayExtractor(
|
||||
model, processor, config, fixed_grid_av=fixed_grid_av
|
||||
)
|
||||
result = extractor.extract(
|
||||
video_path=args.video,
|
||||
audio_path=args.audio,
|
||||
start=args.start,
|
||||
end=end,
|
||||
use_audio_in_video=args.use_audio_in_video,
|
||||
verify_every=args.verify_every,
|
||||
verify_continuous_every=args.verify_continuous_every,
|
||||
show_progress=not args.no_progress,
|
||||
)
|
||||
else:
|
||||
extractor = OpenStreamReplayExtractor(
|
||||
@@ -1811,6 +1713,7 @@ def main() -> None:
|
||||
use_audio_in_video=args.use_audio_in_video,
|
||||
verify_every=args.verify_every,
|
||||
verify_continuous_every=args.verify_continuous_every,
|
||||
show_progress=not args.no_progress,
|
||||
)
|
||||
|
||||
payload: dict[str, Any] = {
|
||||
@@ -1819,11 +1722,7 @@ def main() -> None:
|
||||
"stream_representation": (
|
||||
"single_causal_window"
|
||||
if args.mode == "causal-prefix"
|
||||
else (
|
||||
"sequence_of_2s_media_turns"
|
||||
if args.mode == "commit-replay"
|
||||
else "single_open_media_turn_with_2s_payload_commits"
|
||||
)
|
||||
else "single_open_media_turn_with_2s_payload_commits"
|
||||
),
|
||||
"config": asdict(config),
|
||||
"start": args.start,
|
||||
@@ -1835,7 +1734,7 @@ def main() -> None:
|
||||
"session_metadata": session_metadata,
|
||||
"input_video_fps": input_video_fps,
|
||||
"video_sampling_grid": (
|
||||
"explicit_fixed_past_frames_with_duplication"
|
||||
"explicit_fixed_half_open_past_frames_with_tail_duplication"
|
||||
if fixed_grid_av is not None
|
||||
else "qwen_utility_adaptive_sampling"
|
||||
),
|
||||
|
||||
913
streaming_turn_model_directly.py
Normal file
913
streaming_turn_model_directly.py
Normal file
@@ -0,0 +1,913 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Extract causal Qwen2.5-Omni states directly from original media.
|
||||
|
||||
Unlike ``prepare_easycom_session.py`` and ``prepare_ego4d_2fps.py``, this
|
||||
entry point does not create a re-encoded fixed-frame-rate MP4. It builds a
|
||||
global visual sampling lattice in Python, selects the newest decoded source
|
||||
frame whose presentation timestamp is not later than each lattice time, and
|
||||
passes explicit PIL frame lists plus exact 16-kHz audio slices to Qwen.
|
||||
|
||||
The default mode follows the released Qwen2.5-Omni streaming design as closely
|
||||
as the public Transformers implementation permits:
|
||||
|
||||
* 2 FPS causal visual sampling (configurable for ablations);
|
||||
* two-frame temporal patches, with the last observed frame duplicated when an
|
||||
incomplete input contains an odd number of frames;
|
||||
* a single open media turn with completed 2-second AV payload commits;
|
||||
* persistent Thinker KV for the complete extraction (``--memory-seconds 0``);
|
||||
* one temporary partial-block branch and readout every 200 ms.
|
||||
|
||||
Examples
|
||||
--------
|
||||
Ego4D, directly from one original clip::
|
||||
|
||||
python streaming_turn_model_directly.py \
|
||||
--model-id /data/wangx/Models/Qwen2.5-Omni-3B \
|
||||
--video /data/wangx/Datasets/ego4d/v1/clips/UID.mp4 \
|
||||
--target /data/wangx/Datasets/ego4d/ego4d_target_frame/UID.npy \
|
||||
--mode open-stream-replay \
|
||||
--memory-seconds 20 \
|
||||
--video-fps 2 \
|
||||
--flash-attention-2 \
|
||||
--device cuda:0 \
|
||||
--output features/ego4d_direct/UID.pt
|
||||
|
||||
EasyCom, directly from the retained one-minute source clips::
|
||||
|
||||
python streaming_turn_model_directly.py \
|
||||
--model-id /path/to/Qwen2.5-Omni-3B \
|
||||
--easycom-session-dir /path/to/Main/Video_Compressed/Session_1 \
|
||||
--target target_frame/Session_1.npy \
|
||||
--mode open-stream-replay \
|
||||
--memory-seconds 20 \
|
||||
--video-fps 2 \
|
||||
--flash-attention-2 \
|
||||
--device cuda:0 \
|
||||
--output features/easycom_direct/Session_1.pt
|
||||
|
||||
``--media-only`` performs source/timestamp/audio validation without loading
|
||||
Qwen. This is useful before launching a full extraction on a remote machine.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import shutil
|
||||
import subprocess
|
||||
from collections import OrderedDict
|
||||
from dataclasses import asdict, dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
import streaming_turn_model as core
|
||||
from easycom_manifest import parse_timecode
|
||||
|
||||
|
||||
EPSILON = 1e-7
|
||||
SAMPLE_RATE = 16_000
|
||||
EASYCOM_WIDTH = 1920
|
||||
EASYCOM_HEIGHT = 1080
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class SourceClipSpec:
|
||||
"""One source file in a dense, gap-free media timeline."""
|
||||
|
||||
path: Path
|
||||
duration: float | None = None
|
||||
crop: tuple[int, int, int, int] | None = None
|
||||
crop_rightmost_width: int | None = None
|
||||
audio_path: Path | None = None
|
||||
|
||||
|
||||
@dataclass
|
||||
class IndexedSourceClip:
|
||||
"""Decoded indexing information for one source clip."""
|
||||
|
||||
path: Path
|
||||
audio_path: Path
|
||||
dense_start: float
|
||||
dense_end: float
|
||||
duration: float
|
||||
crop: tuple[int, int, int, int] | None
|
||||
reader: Any
|
||||
source_frame_starts: np.ndarray
|
||||
source_frame_ends: np.ndarray
|
||||
source_fps: float
|
||||
width: int
|
||||
height: int
|
||||
|
||||
@property
|
||||
def num_frames(self) -> int:
|
||||
return int(len(self.source_frame_starts))
|
||||
|
||||
|
||||
def _grid_indices(start: float, end: float, fps: float) -> list[int]:
|
||||
"""Indices of global samples ``n/fps`` in the half-open interval."""
|
||||
|
||||
first = max(0, int(math.ceil(start * fps - EPSILON)))
|
||||
last_exclusive = int(math.ceil(end * fps - EPSILON))
|
||||
if first >= last_exclusive:
|
||||
# Qwen's visual tower needs at least one observed image before its
|
||||
# temporal-patch duplication. Reuse the latest global tick strictly
|
||||
# before ``end``; this is identical to FixedGridAV's short-tail rule.
|
||||
return [max(0, last_exclusive - 1)]
|
||||
return list(range(first, last_exclusive))
|
||||
|
||||
|
||||
def _timestamp_array(reader: Any) -> tuple[np.ndarray, np.ndarray]:
|
||||
"""Return monotonically increasing, zero-origin frame PTS intervals."""
|
||||
|
||||
# Decord releases differ here: some accept an ndarray, while older ones
|
||||
# require a Python list. A list works with both APIs.
|
||||
indices = list(range(len(reader)))
|
||||
raw = reader.get_frame_timestamp(indices)
|
||||
if hasattr(raw, "asnumpy"):
|
||||
raw = raw.asnumpy()
|
||||
timestamps = np.asarray(raw, dtype=np.float64)
|
||||
if timestamps.ndim == 1:
|
||||
starts = timestamps
|
||||
nominal = 1.0 / float(reader.get_avg_fps())
|
||||
ends = np.concatenate((starts[1:], starts[-1:] + nominal))
|
||||
elif timestamps.ndim == 2 and timestamps.shape[1] >= 2:
|
||||
starts = timestamps[:, 0]
|
||||
ends = timestamps[:, 1]
|
||||
else:
|
||||
raise RuntimeError(
|
||||
f"Unexpected decord frame-timestamp shape: {timestamps.shape}"
|
||||
)
|
||||
if not len(starts):
|
||||
raise RuntimeError("Video contains no decoded frames")
|
||||
origin = float(starts[0])
|
||||
starts = starts - origin
|
||||
ends = ends - origin
|
||||
if np.any(np.diff(starts) < -1e-9):
|
||||
raise RuntimeError("Decoded video frame PTS values are not monotonic")
|
||||
return starts, ends
|
||||
|
||||
|
||||
def _decode_audio_exact(path: Path, duration: float) -> np.ndarray:
|
||||
"""Decode, then pad/trim one clip to its dense video duration."""
|
||||
|
||||
completed = subprocess.run(
|
||||
[
|
||||
"ffmpeg",
|
||||
"-nostdin",
|
||||
"-v",
|
||||
"error",
|
||||
"-i",
|
||||
str(path),
|
||||
"-map",
|
||||
"0:a:0",
|
||||
"-vn",
|
||||
"-ac",
|
||||
"1",
|
||||
"-ar",
|
||||
str(SAMPLE_RATE),
|
||||
"-c:a",
|
||||
"pcm_f32le",
|
||||
"-f",
|
||||
"f32le",
|
||||
"pipe:1",
|
||||
],
|
||||
check=False,
|
||||
capture_output=True,
|
||||
)
|
||||
if completed.returncode != 0:
|
||||
message = completed.stderr.decode("utf-8", errors="replace").strip()
|
||||
raise RuntimeError(f"Failed to decode audio from {path}: {message}")
|
||||
|
||||
decoded = np.frombuffer(completed.stdout, dtype="<f4")
|
||||
expected = int(round(duration * SAMPLE_RATE))
|
||||
output = np.zeros(expected, dtype=np.float32)
|
||||
used = min(expected, len(decoded))
|
||||
if used:
|
||||
output[:used] = decoded[:used]
|
||||
return output
|
||||
|
||||
|
||||
class CausalSourceAV:
|
||||
"""Direct PTS-aware reader implementing the existing extractor AV API."""
|
||||
|
||||
sample_rate = SAMPLE_RATE
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
specs: list[SourceClipSpec],
|
||||
*,
|
||||
output_fps: float,
|
||||
required_duration: float,
|
||||
frame_cache_size: int = 8,
|
||||
) -> None:
|
||||
if not specs:
|
||||
raise ValueError("At least one source clip is required")
|
||||
if output_fps <= 0:
|
||||
raise ValueError("output_fps must be positive")
|
||||
if required_duration <= 0:
|
||||
raise ValueError("required_duration must be positive")
|
||||
if frame_cache_size < 0:
|
||||
raise ValueError("frame_cache_size must be non-negative")
|
||||
if shutil.which("ffmpeg") is None:
|
||||
raise RuntimeError("ffmpeg is required to decode source audio")
|
||||
|
||||
try:
|
||||
import decord
|
||||
except ImportError as exc:
|
||||
raise RuntimeError(
|
||||
"Direct extraction requires decord; install "
|
||||
"qwen-omni-utils[decord]"
|
||||
) from exc
|
||||
|
||||
self.fps = float(output_fps)
|
||||
self.required_duration = float(required_duration)
|
||||
self.frame_cache_size = int(frame_cache_size)
|
||||
self._frame_cache: OrderedDict[tuple[int, int], Any] = OrderedDict()
|
||||
self._request_count = 0
|
||||
self._cache_hits = 0
|
||||
self._odd_frame_duplications = 0
|
||||
self._short_interval_fallbacks = 0
|
||||
self._tail_frame_reuses = 0
|
||||
|
||||
indexed: list[IndexedSourceClip] = []
|
||||
audio_parts: list[np.ndarray] = []
|
||||
dense_cursor = 0.0
|
||||
|
||||
for spec_index, spec in enumerate(specs):
|
||||
path = spec.path.expanduser().resolve()
|
||||
if not path.is_file():
|
||||
raise FileNotFoundError(f"Source video does not exist: {path}")
|
||||
print(
|
||||
f"direct-media: indexing [{spec_index + 1}/{len(specs)}] "
|
||||
f"{path.name}",
|
||||
flush=True,
|
||||
)
|
||||
reader = decord.VideoReader(str(path))
|
||||
if not len(reader):
|
||||
raise RuntimeError(f"No video frames decoded from {path}")
|
||||
starts, ends = _timestamp_array(reader)
|
||||
source_fps = float(reader.get_avg_fps())
|
||||
first_shape = tuple(int(value) for value in reader[0].shape)
|
||||
if len(first_shape) != 3:
|
||||
raise RuntimeError(
|
||||
f"Unexpected decoded frame shape {first_shape}: {path}"
|
||||
)
|
||||
height, width = first_shape[:2]
|
||||
|
||||
inferred_duration = len(reader) / source_fps
|
||||
duration = (
|
||||
float(spec.duration)
|
||||
if spec.duration is not None
|
||||
else float(inferred_duration)
|
||||
)
|
||||
if duration <= 0:
|
||||
raise RuntimeError(f"Non-positive source duration for {path}")
|
||||
|
||||
crop = spec.crop
|
||||
if crop is None and spec.crop_rightmost_width is not None:
|
||||
keep_width = int(spec.crop_rightmost_width)
|
||||
if width < keep_width:
|
||||
raise RuntimeError(
|
||||
f"Cannot keep rightmost {keep_width}px from {width}px: {path}"
|
||||
)
|
||||
crop = (width - keep_width, 0, keep_width, height)
|
||||
if crop is not None:
|
||||
x, y, crop_width, crop_height = crop
|
||||
if (
|
||||
x < 0
|
||||
or y < 0
|
||||
or crop_width <= 0
|
||||
or crop_height <= 0
|
||||
or x + crop_width > width
|
||||
or y + crop_height > height
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Invalid crop {crop} for {width}x{height}: {path}"
|
||||
)
|
||||
|
||||
audio_path = (spec.audio_path or path).expanduser().resolve()
|
||||
if not audio_path.is_file():
|
||||
raise FileNotFoundError(f"Audio source does not exist: {audio_path}")
|
||||
|
||||
clip = IndexedSourceClip(
|
||||
path=path,
|
||||
audio_path=audio_path,
|
||||
dense_start=dense_cursor,
|
||||
dense_end=dense_cursor + duration,
|
||||
duration=duration,
|
||||
crop=crop,
|
||||
reader=reader,
|
||||
source_frame_starts=starts,
|
||||
source_frame_ends=ends,
|
||||
source_fps=source_fps,
|
||||
width=width,
|
||||
height=height,
|
||||
)
|
||||
indexed.append(clip)
|
||||
print(
|
||||
f"direct-media: decoding audio [{spec_index + 1}/{len(specs)}] "
|
||||
f"{audio_path.name}",
|
||||
flush=True,
|
||||
)
|
||||
audio_parts.append(_decode_audio_exact(audio_path, duration))
|
||||
dense_cursor = clip.dense_end
|
||||
|
||||
# No later source can be observed before the requested endpoint.
|
||||
if dense_cursor >= required_duration - EPSILON:
|
||||
break
|
||||
|
||||
self.clips = indexed
|
||||
self.dense_starts = np.asarray(
|
||||
[clip.dense_start for clip in indexed], dtype=np.float64
|
||||
)
|
||||
self.dense_duration = float(dense_cursor)
|
||||
if self.dense_duration + EPSILON < self.required_duration:
|
||||
raise RuntimeError(
|
||||
f"Dense source duration is {self.dense_duration:.6f}s, shorter "
|
||||
f"than required {self.required_duration:.6f}s"
|
||||
)
|
||||
self.audio = np.concatenate(audio_parts)
|
||||
required_samples = int(round(required_duration * SAMPLE_RATE))
|
||||
if len(self.audio) < required_samples:
|
||||
self.audio = np.pad(self.audio, (0, required_samples - len(self.audio)))
|
||||
elif len(self.audio) > required_samples:
|
||||
self.audio = self.audio[:required_samples]
|
||||
|
||||
max_grid_index = int(
|
||||
math.ceil(required_duration * self.fps - EPSILON)
|
||||
) - 1
|
||||
self.expected_grid_frames = max(0, max_grid_index + 1)
|
||||
if max_grid_index >= 0:
|
||||
self._source_frame_for_grid(max_grid_index, count_tail_reuse=False)
|
||||
|
||||
def _clip_for_time(self, global_time: float) -> tuple[int, IndexedSourceClip, float]:
|
||||
clip_index = int(np.searchsorted(self.dense_starts, global_time, side="right") - 1)
|
||||
clip_index = max(0, min(clip_index, len(self.clips) - 1))
|
||||
clip = self.clips[clip_index]
|
||||
local_time = global_time - clip.dense_start
|
||||
if local_time < 0:
|
||||
local_time = 0.0
|
||||
return clip_index, clip, local_time
|
||||
|
||||
def _source_frame_for_grid(
|
||||
self,
|
||||
grid_index: int,
|
||||
*,
|
||||
count_tail_reuse: bool = True,
|
||||
) -> tuple[int, int, float, float]:
|
||||
"""Map one global output tick to a source frame at or before it."""
|
||||
|
||||
global_time = grid_index / self.fps
|
||||
clip_index, clip, local_time = self._clip_for_time(global_time)
|
||||
frame_index = int(
|
||||
np.searchsorted(clip.source_frame_starts, local_time + EPSILON, side="right")
|
||||
- 1
|
||||
)
|
||||
frame_index = max(0, min(frame_index, clip.num_frames - 1))
|
||||
source_time = float(clip.source_frame_starts[frame_index])
|
||||
if source_time > local_time + 2 * EPSILON:
|
||||
raise RuntimeError(
|
||||
"Causal frame-selection violation: "
|
||||
f"source PTS {source_time:.9f}s > grid-local time {local_time:.9f}s"
|
||||
)
|
||||
if count_tail_reuse and local_time > float(clip.source_frame_starts[-1]) + EPSILON:
|
||||
self._tail_frame_reuses += 1
|
||||
return clip_index, frame_index, global_time, source_time
|
||||
|
||||
def _load_frame(self, clip_index: int, frame_index: int) -> Any:
|
||||
from PIL import Image
|
||||
|
||||
key = (clip_index, frame_index)
|
||||
cached = self._frame_cache.get(key)
|
||||
if cached is not None:
|
||||
self._cache_hits += 1
|
||||
self._frame_cache.move_to_end(key)
|
||||
return cached.copy()
|
||||
|
||||
clip = self.clips[clip_index]
|
||||
array = clip.reader[frame_index].asnumpy()
|
||||
if clip.crop is not None:
|
||||
x, y, width, height = clip.crop
|
||||
array = array[y : y + height, x : x + width]
|
||||
image = Image.fromarray(np.asarray(array)).copy()
|
||||
if self.frame_cache_size:
|
||||
self._frame_cache[key] = image
|
||||
self._frame_cache.move_to_end(key)
|
||||
while len(self._frame_cache) > self.frame_cache_size:
|
||||
self._frame_cache.popitem(last=False)
|
||||
return image.copy()
|
||||
|
||||
def frames(self, start: float, end: float) -> list[Any]:
|
||||
"""Return causal global-grid frames in ``[start, end)``."""
|
||||
|
||||
if end <= start:
|
||||
raise ValueError(f"Empty video interval [{start}, {end}]")
|
||||
raw_first = int(math.ceil(start * self.fps - EPSILON))
|
||||
raw_last_exclusive = int(math.ceil(end * self.fps - EPSILON))
|
||||
grid_indices = _grid_indices(start, end, self.fps)
|
||||
if raw_first >= raw_last_exclusive:
|
||||
self._short_interval_fallbacks += 1
|
||||
|
||||
output = []
|
||||
for grid_index in grid_indices:
|
||||
clip_index, frame_index, _, _ = self._source_frame_for_grid(grid_index)
|
||||
output.append(self._load_frame(clip_index, frame_index))
|
||||
self._request_count += 1
|
||||
if len(output) % 2:
|
||||
output.append(output[-1].copy())
|
||||
self._odd_frame_duplications += 1
|
||||
return output
|
||||
|
||||
def audio_slice(self, start: float, end: float) -> np.ndarray:
|
||||
"""Return the exact 16-kHz mono samples in ``[start, end)``."""
|
||||
|
||||
first = max(0, int(math.floor(start * SAMPLE_RATE + EPSILON)))
|
||||
last = min(
|
||||
len(self.audio), int(math.floor(end * SAMPLE_RATE + EPSILON))
|
||||
)
|
||||
if last <= first:
|
||||
raise ValueError(f"Empty audio interval [{start}, {end}]")
|
||||
return self.audio[first:last]
|
||||
|
||||
def sample_mapping(self, grid_index: int) -> dict[str, Any]:
|
||||
"""Describe one visual-grid mapping without decoding its pixels."""
|
||||
|
||||
clip_index, frame_index, grid_time, source_time = self._source_frame_for_grid(
|
||||
grid_index, count_tail_reuse=False
|
||||
)
|
||||
clip = self.clips[clip_index]
|
||||
return {
|
||||
"grid_index": int(grid_index),
|
||||
"grid_time": float(grid_time),
|
||||
"source_clip": str(clip.path),
|
||||
"source_frame_index": int(frame_index),
|
||||
"source_frame_pts_seconds": float(source_time),
|
||||
"dense_source_frame_time": float(clip.dense_start + source_time),
|
||||
}
|
||||
|
||||
def metadata(self) -> dict[str, Any]:
|
||||
return {
|
||||
"decoder": "decord_with_frame_pts_and_ffmpeg_pcm_audio",
|
||||
"visual_sampling": "global_grid_latest_source_pts_not_after_tick",
|
||||
"interval_convention": "half_open_[start,end)",
|
||||
"output_video_fps": self.fps,
|
||||
"expected_grid_frames": self.expected_grid_frames,
|
||||
"required_duration_seconds": self.required_duration,
|
||||
"dense_source_duration_seconds": self.dense_duration,
|
||||
"audio_sample_rate": SAMPLE_RATE,
|
||||
"audio_samples": int(len(self.audio)),
|
||||
"temporal_patch_padding": "duplicate_last_observed_frame_when_odd",
|
||||
"source_clips": [
|
||||
{
|
||||
"path": str(clip.path),
|
||||
"audio_path": str(clip.audio_path),
|
||||
"dense_start": clip.dense_start,
|
||||
"dense_end": clip.dense_end,
|
||||
"duration": clip.duration,
|
||||
"source_fps": clip.source_fps,
|
||||
"source_frames": clip.num_frames,
|
||||
"source_size": [clip.width, clip.height],
|
||||
"crop": list(clip.crop) if clip.crop is not None else None,
|
||||
"first_source_pts": float(clip.source_frame_starts[0]),
|
||||
"last_source_pts": float(clip.source_frame_starts[-1]),
|
||||
}
|
||||
for clip in self.clips
|
||||
],
|
||||
"runtime_counters": {
|
||||
"frame_requests": self._request_count,
|
||||
"frame_cache_hits": self._cache_hits,
|
||||
"odd_frame_duplications": self._odd_frame_duplications,
|
||||
"short_interval_fallbacks": self._short_interval_fallbacks,
|
||||
"tail_frame_reuses": self._tail_frame_reuses,
|
||||
},
|
||||
}
|
||||
|
||||
|
||||
def _manifest_session_specs(
|
||||
manifest_path: Path,
|
||||
session_name: str,
|
||||
) -> tuple[list[SourceClipSpec], dict[str, Any], Path | None]:
|
||||
payload = json.loads(manifest_path.read_text())
|
||||
matches = [
|
||||
session
|
||||
for session in payload.get("sessions", [])
|
||||
if session.get("session") == session_name
|
||||
]
|
||||
if len(matches) != 1:
|
||||
raise RuntimeError(
|
||||
f"Expected exactly one {session_name!r} in {manifest_path}, "
|
||||
f"found {len(matches)}"
|
||||
)
|
||||
session = matches[0]
|
||||
specs = []
|
||||
for record in session.get("clips", []):
|
||||
duration = record.get("video_duration_seconds")
|
||||
crop_record = record.get("egospeak_crop")
|
||||
crop = None
|
||||
if crop_record is not None:
|
||||
crop = (
|
||||
int(crop_record["x"]),
|
||||
int(crop_record["y"]),
|
||||
int(crop_record["width"]),
|
||||
int(crop_record["height"]),
|
||||
)
|
||||
specs.append(
|
||||
SourceClipSpec(
|
||||
path=Path(record["path"]),
|
||||
duration=float(duration) if duration is not None else None,
|
||||
crop=crop,
|
||||
)
|
||||
)
|
||||
if not specs:
|
||||
raise RuntimeError(f"No source clips recorded for {session_name}")
|
||||
target_record = session.get("target") or {}
|
||||
target_path = Path(target_record["path"]) if target_record.get("path") else None
|
||||
return specs, session, target_path
|
||||
|
||||
|
||||
def _easycom_directory_specs(session_dir: Path) -> list[SourceClipSpec]:
|
||||
session_dir = session_dir.expanduser().resolve()
|
||||
if not session_dir.is_dir():
|
||||
raise FileNotFoundError(f"EasyCom session directory does not exist: {session_dir}")
|
||||
paths = sorted(session_dir.glob("*.mp4"), key=lambda path: parse_timecode(path.stem))
|
||||
if not paths:
|
||||
raise RuntimeError(f"No MP4 clips found in {session_dir}")
|
||||
return [
|
||||
SourceClipSpec(path=path, crop_rightmost_width=EASYCOM_WIDTH)
|
||||
for path in paths
|
||||
]
|
||||
|
||||
|
||||
def _parse_crop(value: str | None) -> tuple[int, int, int, int] | None:
|
||||
if value is None:
|
||||
return None
|
||||
try:
|
||||
values = tuple(int(item) for item in value.split(","))
|
||||
except ValueError as exc:
|
||||
raise argparse.ArgumentTypeError("crop must be x,y,width,height") from exc
|
||||
if len(values) != 4:
|
||||
raise argparse.ArgumentTypeError("crop must be x,y,width,height")
|
||||
return values
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
source = parser.add_mutually_exclusive_group(required=True)
|
||||
source.add_argument("--video", type=Path, help="One original video clip")
|
||||
source.add_argument(
|
||||
"--easycom-session-dir",
|
||||
type=Path,
|
||||
help="Original EasyCom Session_N directory containing retained MP4 clips",
|
||||
)
|
||||
source.add_argument(
|
||||
"--easycom-manifest",
|
||||
type=Path,
|
||||
help="easycom_manifest.py JSON; also requires --session",
|
||||
)
|
||||
parser.add_argument("--session", help="Session_N selected from --easycom-manifest")
|
||||
parser.add_argument(
|
||||
"--audio",
|
||||
type=Path,
|
||||
help="Optional separate audio for --video; defaults to its embedded track",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--crop",
|
||||
help="Optional crop for --video as x,y,width,height",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--use-audio-in-video",
|
||||
action="store_true",
|
||||
help="Accepted for command compatibility; direct extraction always uses AV interleaving",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--target",
|
||||
type=Path,
|
||||
help="target_perframe [T,3]; T is the authoritative full duration",
|
||||
)
|
||||
parser.add_argument("--output", type=Path)
|
||||
parser.add_argument(
|
||||
"--mode",
|
||||
choices=("causal-prefix", "open-stream-replay"),
|
||||
default="open-stream-replay",
|
||||
)
|
||||
parser.add_argument("--model-id", default=core.MODEL_ID)
|
||||
parser.add_argument("--start", type=float, default=0.0)
|
||||
parser.add_argument("--end", type=float)
|
||||
parser.add_argument("--step-seconds", type=float, default=0.2)
|
||||
parser.add_argument("--native-block-seconds", type=float, default=2.0)
|
||||
parser.add_argument(
|
||||
"--memory-seconds",
|
||||
type=float,
|
||||
default=0.0,
|
||||
help="0 keeps Thinker KV for the full extraction; positive values use a rolling window",
|
||||
)
|
||||
parser.add_argument("--video-fps", type=float, default=2.0)
|
||||
parser.add_argument("--video-min-pixels", type=int, default=64 * 28 * 28)
|
||||
parser.add_argument("--video-max-pixels", type=int, default=128 * 28 * 28)
|
||||
parser.add_argument("--frame-cache-size", type=int, default=8)
|
||||
parser.add_argument("--verify-every", type=int, default=0)
|
||||
parser.add_argument("--verify-continuous-every", type=int, default=0)
|
||||
parser.add_argument("--no-progress", action="store_true")
|
||||
parser.add_argument("--head-checkpoint", type=Path)
|
||||
parser.add_argument(
|
||||
"--dtype",
|
||||
choices=("auto", "bfloat16", "float16", "float32"),
|
||||
default="bfloat16",
|
||||
)
|
||||
parser.add_argument("--flash-attention-2", action="store_true")
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument(
|
||||
"--media-only",
|
||||
action="store_true",
|
||||
help="Validate/index direct media and exit before loading Qwen",
|
||||
)
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _resolve_duration_and_target(
|
||||
args: argparse.Namespace,
|
||||
manifest_target: Path | None,
|
||||
) -> tuple[Path | None, int | None, float]:
|
||||
target_path = args.target
|
||||
if target_path is None and manifest_target is not None and manifest_target.is_file():
|
||||
target_path = manifest_target
|
||||
target_steps = None
|
||||
if target_path is not None:
|
||||
target_path = target_path.expanduser().resolve()
|
||||
target_steps = core.target_num_steps(target_path)
|
||||
if not math.isclose(args.start, 0.0, rel_tol=0.0, abs_tol=1e-9):
|
||||
raise SystemExit("Full target extraction requires --start 0")
|
||||
target_end = target_steps * args.step_seconds
|
||||
if args.end is not None and not math.isclose(
|
||||
args.end, target_end, rel_tol=0.0, abs_tol=1e-6
|
||||
):
|
||||
raise SystemExit(
|
||||
f"--end {args.end} conflicts with target length {target_steps} "
|
||||
f"at {args.step_seconds}s ({target_end}s)"
|
||||
)
|
||||
end = target_end
|
||||
else:
|
||||
if args.end is None:
|
||||
raise SystemExit("Pass --target for a full run or --end for a bounded test")
|
||||
end = float(args.end)
|
||||
if end <= args.start:
|
||||
raise SystemExit("--end must be greater than --start")
|
||||
return target_path, target_steps, end
|
||||
|
||||
|
||||
def _effective_memory_seconds(
|
||||
requested: float,
|
||||
*,
|
||||
duration: float,
|
||||
block_seconds: float,
|
||||
) -> float:
|
||||
if requested < 0:
|
||||
raise SystemExit("--memory-seconds must be >= 0")
|
||||
if requested > 0:
|
||||
return requested
|
||||
blocks = max(1, int(math.ceil(duration / block_seconds - EPSILON)))
|
||||
return blocks * block_seconds
|
||||
|
||||
|
||||
def _print_media_summary(reader: CausalSourceAV, start: float, end: float) -> None:
|
||||
print(
|
||||
f"Direct media valid: {len(reader.clips)} source clip(s), "
|
||||
f"required={end - start:.3f}s, grid={reader.expected_grid_frames} "
|
||||
f"frames at {reader.fps:g} FPS, audio={len(reader.audio)} samples"
|
||||
)
|
||||
candidates = [0]
|
||||
if reader.expected_grid_frames > 1:
|
||||
candidates.append(reader.expected_grid_frames - 1)
|
||||
for grid_index in candidates:
|
||||
mapping = reader.sample_mapping(grid_index)
|
||||
print(
|
||||
" grid "
|
||||
f"{mapping['grid_time']:.3f}s -> "
|
||||
f"{Path(mapping['source_clip']).name} "
|
||||
f"frame={mapping['source_frame_index']} "
|
||||
f"pts={mapping['source_frame_pts_seconds']:.6f}s"
|
||||
)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
core._configure_qwen_logging()
|
||||
if args.verify_every < 0 or args.verify_continuous_every < 0:
|
||||
raise SystemExit("verification intervals must be >= 0")
|
||||
if args.frame_cache_size < 0:
|
||||
raise SystemExit("--frame-cache-size must be >= 0")
|
||||
if args.output is None and not args.media_only:
|
||||
raise SystemExit("--output is required unless --media-only is used")
|
||||
if args.audio is not None and args.video is None:
|
||||
raise SystemExit("--audio is supported only together with --video")
|
||||
if args.session is not None and args.easycom_manifest is None:
|
||||
raise SystemExit("--session is used only with --easycom-manifest")
|
||||
if args.easycom_manifest is not None and args.session is None:
|
||||
raise SystemExit("--easycom-manifest requires --session Session_N")
|
||||
|
||||
manifest_record = None
|
||||
manifest_target = None
|
||||
if args.video is not None:
|
||||
specs = [
|
||||
SourceClipSpec(
|
||||
path=args.video,
|
||||
crop=_parse_crop(args.crop),
|
||||
audio_path=args.audio,
|
||||
)
|
||||
]
|
||||
source_kind = "single_original_video"
|
||||
elif args.easycom_session_dir is not None:
|
||||
if args.crop is not None:
|
||||
raise SystemExit("EasyCom direct extraction uses the fixed rightmost-1920 crop")
|
||||
specs = _easycom_directory_specs(args.easycom_session_dir)
|
||||
source_kind = "easycom_dense_original_clips"
|
||||
else:
|
||||
if args.crop is not None:
|
||||
raise SystemExit("EasyCom manifest supplies its crop metadata")
|
||||
specs, manifest_record, manifest_target = _manifest_session_specs(
|
||||
args.easycom_manifest.expanduser().resolve(), args.session
|
||||
)
|
||||
source_kind = "easycom_manifest_original_clips"
|
||||
|
||||
target_path, target_steps, end = _resolve_duration_and_target(
|
||||
args, manifest_target
|
||||
)
|
||||
required_duration = end
|
||||
reader = CausalSourceAV(
|
||||
specs,
|
||||
output_fps=args.video_fps,
|
||||
required_duration=required_duration,
|
||||
frame_cache_size=args.frame_cache_size,
|
||||
)
|
||||
is_easycom = (
|
||||
args.easycom_session_dir is not None
|
||||
or args.easycom_manifest is not None
|
||||
)
|
||||
if is_easycom:
|
||||
for clip in reader.clips:
|
||||
if clip.height != EASYCOM_HEIGHT:
|
||||
raise RuntimeError(
|
||||
f"Expected EasyCom height {EASYCOM_HEIGHT}, got {clip.height}: "
|
||||
f"{clip.path}"
|
||||
)
|
||||
if not math.isclose(
|
||||
clip.source_fps, 20.0, rel_tol=0.0, abs_tol=1e-6
|
||||
):
|
||||
raise RuntimeError(
|
||||
f"Expected EasyCom source FPS 20, got {clip.source_fps:g}: "
|
||||
f"{clip.path}"
|
||||
)
|
||||
if target_steps is not None:
|
||||
source_frames = sum(clip.num_frames for clip in reader.clips)
|
||||
media_steps = source_frames // 4
|
||||
if media_steps != target_steps:
|
||||
raise RuntimeError(
|
||||
f"EasyCom dense source gives {media_steps} complete 5-Hz "
|
||||
f"steps from {source_frames} frames, but target has "
|
||||
f"{target_steps}"
|
||||
)
|
||||
_print_media_summary(reader, args.start, end)
|
||||
if args.media_only:
|
||||
return
|
||||
|
||||
duration = end - args.start
|
||||
effective_memory = _effective_memory_seconds(
|
||||
args.memory_seconds,
|
||||
duration=duration,
|
||||
block_seconds=args.native_block_seconds,
|
||||
)
|
||||
if args.memory_seconds == 0 and duration > 120:
|
||||
print(
|
||||
"WARNING: persistent Thinker KV grows for the entire media "
|
||||
f"duration ({duration:.1f}s). Long EasyCom sessions can exceed "
|
||||
"the model context or GPU memory; a positive --memory-seconds is "
|
||||
"a bounded-memory experiment rather than native persistent-KV "
|
||||
"semantics.",
|
||||
flush=True,
|
||||
)
|
||||
config = core.StreamConfig(
|
||||
step_seconds=args.step_seconds,
|
||||
native_block_seconds=args.native_block_seconds,
|
||||
memory_seconds=effective_memory,
|
||||
video_fps=args.video_fps,
|
||||
video_min_pixels=args.video_min_pixels,
|
||||
video_max_pixels=args.video_max_pixels,
|
||||
)
|
||||
config.validate()
|
||||
|
||||
model, processor = core.load_qwen_omni(
|
||||
args.model_id,
|
||||
dtype=args.dtype,
|
||||
device_map={"": args.device},
|
||||
use_flash_attention_2=args.flash_attention_2,
|
||||
)
|
||||
model.requires_grad_(False)
|
||||
media_path_for_api = reader.clips[0].path
|
||||
|
||||
if args.mode == "causal-prefix":
|
||||
extractor = core.CausalPrefixExtractor(
|
||||
model, processor, config, fixed_grid_av=reader
|
||||
)
|
||||
result = extractor.extract(
|
||||
video_path=media_path_for_api,
|
||||
audio_path=None,
|
||||
start=args.start,
|
||||
end=end,
|
||||
use_audio_in_video=True,
|
||||
show_progress=not args.no_progress,
|
||||
)
|
||||
else:
|
||||
extractor = core.OpenStreamReplayExtractor(
|
||||
model, processor, config, fixed_grid_av=reader
|
||||
)
|
||||
result = extractor.extract(
|
||||
video_path=media_path_for_api,
|
||||
audio_path=None,
|
||||
start=args.start,
|
||||
end=end,
|
||||
use_audio_in_video=True,
|
||||
verify_every=args.verify_every,
|
||||
verify_continuous_every=args.verify_continuous_every,
|
||||
show_progress=not args.no_progress,
|
||||
)
|
||||
|
||||
if target_steps is not None and len(result.times) != target_steps:
|
||||
raise RuntimeError(
|
||||
f"Extractor returned {len(result.times)} states, expected {target_steps}"
|
||||
)
|
||||
|
||||
persistent_memory = args.memory_seconds == 0
|
||||
payload: dict[str, Any] = {
|
||||
"model_id": args.model_id,
|
||||
"mode": args.mode,
|
||||
"stream_representation": (
|
||||
"single_causal_prefix"
|
||||
if args.mode == "causal-prefix"
|
||||
else "single_open_media_turn_with_2s_payload_commits"
|
||||
),
|
||||
"config": asdict(config),
|
||||
"requested_memory_seconds": args.memory_seconds,
|
||||
"persistent_thinker_kv": persistent_memory,
|
||||
"start": args.start,
|
||||
"end": end,
|
||||
"num_steps": len(result.times),
|
||||
"label_indices": torch.arange(len(result.times), dtype=torch.long),
|
||||
"valid_10_horizon_steps": max(0, len(result.times) - 10),
|
||||
"target_path": str(target_path) if target_path is not None else None,
|
||||
"source_kind": source_kind,
|
||||
"source_manifest_record": manifest_record,
|
||||
"direct_media": reader.metadata(),
|
||||
"video_sampling_grid": (
|
||||
"direct_pts_causal_global_grid_half_open_with_tail_duplication"
|
||||
),
|
||||
"open_stream_policy": (
|
||||
{
|
||||
"commit": "completed_2s_av_payload_only",
|
||||
"temporary": "partial_payload_plus_media_closure_plus_readout",
|
||||
"tmrope_positions_per_chunk": core._tmrope_positions_per_chunk(model),
|
||||
"memory": (
|
||||
"persistent_full_extraction_thinker_kv"
|
||||
if persistent_memory
|
||||
else "rolling_complete_blocks_plus_current_partial_block"
|
||||
),
|
||||
}
|
||||
if args.mode == "open-stream-replay"
|
||||
else None
|
||||
),
|
||||
"times": result.times,
|
||||
"states": result.states,
|
||||
"verification": result.verification,
|
||||
}
|
||||
if args.head_checkpoint is not None:
|
||||
head = core._load_head(
|
||||
args.head_checkpoint, core._hidden_size(model), model.device
|
||||
)
|
||||
with torch.inference_mode():
|
||||
logits = head(
|
||||
result.states.to(model.device, dtype=model.dtype)
|
||||
).float().cpu()
|
||||
payload["logits"] = logits
|
||||
payload["probabilities"] = logits.softmax(dim=-1)
|
||||
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
torch.save(payload, args.output)
|
||||
print(f"Saved {len(result.times)} causal states to {args.output}")
|
||||
print(f"states: {tuple(result.states.shape)}")
|
||||
if result.verification:
|
||||
kinds = sorted({item.get("kind", "unspecified") for item in result.verification})
|
||||
print(f"verification checks: {len(result.verification)}")
|
||||
for kind in kinds:
|
||||
checks = [item for item in result.verification if item.get("kind") == kind]
|
||||
print(
|
||||
f" {kind}: n={len(checks)}, "
|
||||
f"worst max_abs={max(item['max_abs'] for item in checks):.6g}, "
|
||||
f"lowest cosine={min(item['cosine'] for item in checks):.8f}"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
476
train_turn_head.py
Normal file
476
train_turn_head.py
Normal file
@@ -0,0 +1,476 @@
|
||||
"""Train the frozen-Omni EasyCom anticipation head.
|
||||
|
||||
The defaults follow EgoSpeak's EasyCom LSTR optimization config where the
|
||||
architectures are comparable: 50 epochs, Adam, 7e-5 learning rate, 5e-5 weight
|
||||
decay, linear warmup for 10 epochs, cosine decay, dropout 0.2, and MCE loss.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import argparse
|
||||
import json
|
||||
import math
|
||||
import random
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
from torch.optim.lr_scheduler import LambdaLR
|
||||
from torch.utils.data import DataLoader, TensorDataset
|
||||
from tqdm.auto import tqdm
|
||||
|
||||
from easycom_head import (
|
||||
CLASS_NAMES,
|
||||
NUM_CLASSES,
|
||||
NUM_HORIZONS,
|
||||
OFFICIAL_TRAIN_SESSIONS,
|
||||
STEP_SECONDS,
|
||||
TurnPredictionHead,
|
||||
average_precision_metrics,
|
||||
egospeak_multihorizon_loss,
|
||||
load_session_data,
|
||||
miniroad_multihorizon_loss,
|
||||
normalize_session_name,
|
||||
validate_disjoint_sessions,
|
||||
)
|
||||
|
||||
|
||||
def parse_args() -> argparse.Namespace:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--feature-root", type=Path, required=True)
|
||||
parser.add_argument("--target-root", type=Path, required=True)
|
||||
parser.add_argument("--output", type=Path, required=True)
|
||||
parser.add_argument(
|
||||
"--sessions",
|
||||
nargs="+",
|
||||
default=list(OFFICIAL_TRAIN_SESSIONS),
|
||||
help="Training sessions; default is the official nine-session split.",
|
||||
)
|
||||
parser.add_argument(
|
||||
"--dev-sessions",
|
||||
nargs="*",
|
||||
default=[],
|
||||
help=(
|
||||
"Optional held-out training sessions used for checkpoint selection. "
|
||||
"They are automatically removed from --sessions. Never use 1/2/3."
|
||||
),
|
||||
)
|
||||
parser.add_argument("--epochs", type=int, default=50)
|
||||
parser.add_argument("--batch-size", type=int, default=512)
|
||||
parser.add_argument("--learning-rate", type=float, default=7e-5)
|
||||
parser.add_argument("--weight-decay", type=float, default=5e-5)
|
||||
parser.add_argument("--warmup-epochs", type=float, default=10.0)
|
||||
parser.add_argument("--warmup-factor", type=float, default=0.3)
|
||||
parser.add_argument("--dropout", type=float, default=0.2)
|
||||
parser.add_argument(
|
||||
"--loss",
|
||||
choices=("lstr-mce", "miniroad-normalized"),
|
||||
default="lstr-mce",
|
||||
)
|
||||
parser.add_argument("--device", default="cuda:0")
|
||||
parser.add_argument("--num-workers", type=int, default=0)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument(
|
||||
"--expected-mode",
|
||||
default="causal-prefix",
|
||||
help="Expected feature mode, or 'any' to disable the check.",
|
||||
)
|
||||
parser.add_argument("--no-progress", action="store_true")
|
||||
return parser.parse_args()
|
||||
|
||||
|
||||
def _validate_args(args: argparse.Namespace) -> None:
|
||||
if args.epochs <= 0:
|
||||
raise ValueError("--epochs must be positive")
|
||||
if args.batch_size <= 0:
|
||||
raise ValueError("--batch-size must be positive")
|
||||
if args.learning_rate <= 0:
|
||||
raise ValueError("--learning-rate must be positive")
|
||||
if args.weight_decay < 0:
|
||||
raise ValueError("--weight-decay cannot be negative")
|
||||
if not 0 <= args.warmup_epochs <= args.epochs:
|
||||
raise ValueError("--warmup-epochs must be in [0, epochs]")
|
||||
if not 0 < args.warmup_factor <= 1:
|
||||
raise ValueError("--warmup-factor must be in (0, 1]")
|
||||
if not 0 <= args.dropout < 1:
|
||||
raise ValueError("--dropout must be in [0, 1)")
|
||||
if args.num_workers < 0:
|
||||
raise ValueError("--num-workers cannot be negative")
|
||||
|
||||
|
||||
def _resolve_device(name: str) -> torch.device:
|
||||
device = torch.device(name)
|
||||
if device.type == "cuda" and not torch.cuda.is_available():
|
||||
raise RuntimeError(f"CUDA requested but unavailable: {name}")
|
||||
return device
|
||||
|
||||
|
||||
def _seed_everything(seed: int) -> None:
|
||||
random.seed(seed)
|
||||
np.random.seed(seed)
|
||||
torch.manual_seed(seed)
|
||||
if torch.cuda.is_available():
|
||||
torch.cuda.manual_seed_all(seed)
|
||||
|
||||
|
||||
def _load_examples(
|
||||
feature_root: Path,
|
||||
target_root: Path,
|
||||
sessions: list[str],
|
||||
expected_mode: str | None,
|
||||
) -> tuple[torch.Tensor, torch.Tensor, int, list[dict[str, Any]]]:
|
||||
all_states: list[torch.Tensor] = []
|
||||
all_targets: list[torch.Tensor] = []
|
||||
metadata: list[dict[str, Any]] = []
|
||||
hidden_size: int | None = None
|
||||
for name in sessions:
|
||||
data = load_session_data(
|
||||
feature_root,
|
||||
target_root,
|
||||
name,
|
||||
expected_mode=expected_mode,
|
||||
)
|
||||
states, targets, _ = data.future_examples(NUM_HORIZONS)
|
||||
if hidden_size is None:
|
||||
hidden_size = data.hidden_size
|
||||
elif data.hidden_size != hidden_size:
|
||||
raise ValueError(
|
||||
f"Feature dimensions differ: {name} has {data.hidden_size}, "
|
||||
f"expected {hidden_size}"
|
||||
)
|
||||
all_states.append(states)
|
||||
all_targets.append(targets)
|
||||
metadata.append(
|
||||
{
|
||||
"session": name,
|
||||
"num_states": len(data.states),
|
||||
"num_training_examples": len(states),
|
||||
**data.feature_metadata,
|
||||
}
|
||||
)
|
||||
print(
|
||||
f"{name}: states={len(data.states)}, "
|
||||
f"valid_10_horizon_examples={len(states)}"
|
||||
)
|
||||
if hidden_size is None:
|
||||
raise ValueError("No sessions selected")
|
||||
return (
|
||||
torch.cat(all_states, dim=0),
|
||||
torch.cat(all_targets, dim=0),
|
||||
hidden_size,
|
||||
metadata,
|
||||
)
|
||||
|
||||
|
||||
def _scheduler(
|
||||
optimizer: torch.optim.Optimizer,
|
||||
*,
|
||||
steps_per_epoch: int,
|
||||
epochs: int,
|
||||
warmup_epochs: float,
|
||||
warmup_factor: float,
|
||||
) -> LambdaLR:
|
||||
total_steps = max(1, steps_per_epoch * epochs)
|
||||
warmup_steps = steps_per_epoch * warmup_epochs
|
||||
|
||||
def multiplier(step: int) -> float:
|
||||
cosine = 0.5 * (1.0 + math.cos(math.pi * step / total_steps))
|
||||
if warmup_steps > 0 and step < warmup_steps:
|
||||
alpha = step / warmup_steps
|
||||
warmup = warmup_factor * (1.0 - alpha) + alpha
|
||||
else:
|
||||
warmup = 1.0
|
||||
return warmup * cosine
|
||||
|
||||
return LambdaLR(optimizer, multiplier)
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def _evaluate_valid_horizons(
|
||||
head: TurnPredictionHead,
|
||||
states: torch.Tensor,
|
||||
targets: torch.Tensor,
|
||||
device: torch.device,
|
||||
batch_size: int,
|
||||
) -> dict[str, Any]:
|
||||
head.eval()
|
||||
chunks: list[torch.Tensor] = []
|
||||
for start in range(0, len(states), batch_size):
|
||||
x = states[start : start + batch_size].to(
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
non_blocking=True,
|
||||
)
|
||||
chunks.append(head(x).softmax(dim=-1).cpu())
|
||||
probabilities = torch.cat(chunks, dim=0).numpy()
|
||||
target_array = targets.numpy()
|
||||
|
||||
results = []
|
||||
for horizon in range(NUM_HORIZONS):
|
||||
result = average_precision_metrics(
|
||||
probabilities[:, horizon],
|
||||
target_array[:, horizon],
|
||||
)
|
||||
result["horizon"] = horizon + 1
|
||||
result["seconds"] = (horizon + 1) * STEP_SECONDS
|
||||
results.append(result)
|
||||
return {
|
||||
"horizons": results,
|
||||
"mean_map": float(np.mean([item["map"] for item in results])),
|
||||
}
|
||||
|
||||
|
||||
def _checkpoint(
|
||||
head: TurnPredictionHead,
|
||||
*,
|
||||
epoch: int,
|
||||
best_epoch: int,
|
||||
best_dev_map: float | None,
|
||||
train_loss: float,
|
||||
args: argparse.Namespace,
|
||||
train_sessions: list[str],
|
||||
dev_sessions: list[str],
|
||||
hidden_size: int,
|
||||
feature_metadata: list[dict[str, Any]],
|
||||
history: list[dict[str, Any]],
|
||||
) -> dict[str, Any]:
|
||||
return {
|
||||
"format_version": 1,
|
||||
"head": head.state_dict(),
|
||||
"head_config": {
|
||||
"hidden_size": hidden_size,
|
||||
"num_horizons": NUM_HORIZONS,
|
||||
"num_classes": NUM_CLASSES,
|
||||
"dropout": args.dropout,
|
||||
},
|
||||
"alignment": {
|
||||
"state_label_offset": 0,
|
||||
"future_offsets": list(range(1, NUM_HORIZONS + 1)),
|
||||
"future_seconds": [
|
||||
(index + 1) * STEP_SECONDS for index in range(NUM_HORIZONS)
|
||||
],
|
||||
},
|
||||
"classes": list(CLASS_NAMES),
|
||||
"epoch": epoch,
|
||||
"best_epoch": best_epoch,
|
||||
"best_dev_map": best_dev_map,
|
||||
"train_loss": train_loss,
|
||||
"train_sessions": train_sessions,
|
||||
"dev_sessions": dev_sessions,
|
||||
"feature_metadata": feature_metadata,
|
||||
"training_config": {
|
||||
"epochs": args.epochs,
|
||||
"batch_size": args.batch_size,
|
||||
"learning_rate": args.learning_rate,
|
||||
"weight_decay": args.weight_decay,
|
||||
"warmup_epochs": args.warmup_epochs,
|
||||
"warmup_factor": args.warmup_factor,
|
||||
"loss": args.loss,
|
||||
"optimizer": "Adam",
|
||||
"scheduler": "linear_warmup_times_cosine",
|
||||
"seed": args.seed,
|
||||
"expected_mode": args.expected_mode,
|
||||
},
|
||||
"history": history,
|
||||
}
|
||||
|
||||
|
||||
def main() -> None:
|
||||
args = parse_args()
|
||||
_validate_args(args)
|
||||
_seed_everything(args.seed)
|
||||
device = _resolve_device(args.device)
|
||||
expected_mode = None if args.expected_mode == "any" else args.expected_mode
|
||||
|
||||
requested_train = [normalize_session_name(item) for item in args.sessions]
|
||||
dev_sessions = [normalize_session_name(item) for item in args.dev_sessions]
|
||||
train_sessions = [item for item in requested_train if item not in dev_sessions]
|
||||
validate_disjoint_sessions(train_sessions, dev_sessions)
|
||||
if not train_sessions:
|
||||
raise ValueError("No training sessions remain after removing dev sessions")
|
||||
forbidden_test = {"Session_1", "Session_2", "Session_3"}
|
||||
leaked = forbidden_test & (set(train_sessions) | set(dev_sessions))
|
||||
if leaked:
|
||||
raise ValueError(
|
||||
f"Official test sessions cannot be used for training/selection: "
|
||||
f"{sorted(leaked)}"
|
||||
)
|
||||
|
||||
print(f"device: {device}")
|
||||
print(f"train sessions: {', '.join(train_sessions)}")
|
||||
print(f"dev sessions: {', '.join(dev_sessions) if dev_sessions else 'none'}")
|
||||
states, targets, hidden_size, feature_metadata = _load_examples(
|
||||
args.feature_root,
|
||||
args.target_root,
|
||||
train_sessions,
|
||||
expected_mode,
|
||||
)
|
||||
print(f"training examples: {len(states)}, feature_dim={hidden_size}")
|
||||
dev_states: torch.Tensor | None = None
|
||||
dev_targets: torch.Tensor | None = None
|
||||
if dev_sessions:
|
||||
dev_states, dev_targets, dev_hidden_size, _ = _load_examples(
|
||||
args.feature_root,
|
||||
args.target_root,
|
||||
dev_sessions,
|
||||
expected_mode,
|
||||
)
|
||||
if dev_hidden_size != hidden_size:
|
||||
raise ValueError(
|
||||
f"Dev feature dimension {dev_hidden_size} differs from "
|
||||
f"training dimension {hidden_size}"
|
||||
)
|
||||
print(f"development examples: {len(dev_states)}")
|
||||
|
||||
dataset = TensorDataset(states, targets)
|
||||
generator = torch.Generator().manual_seed(args.seed)
|
||||
loader = DataLoader(
|
||||
dataset,
|
||||
batch_size=args.batch_size,
|
||||
shuffle=True,
|
||||
num_workers=args.num_workers,
|
||||
pin_memory=device.type == "cuda",
|
||||
persistent_workers=args.num_workers > 0,
|
||||
generator=generator,
|
||||
)
|
||||
head = TurnPredictionHead(
|
||||
hidden_size=hidden_size,
|
||||
num_horizons=NUM_HORIZONS,
|
||||
num_classes=NUM_CLASSES,
|
||||
dropout=args.dropout,
|
||||
).to(device)
|
||||
optimizer = torch.optim.Adam(
|
||||
head.parameters(),
|
||||
lr=args.learning_rate,
|
||||
weight_decay=args.weight_decay,
|
||||
)
|
||||
scheduler = _scheduler(
|
||||
optimizer,
|
||||
steps_per_epoch=len(loader),
|
||||
epochs=args.epochs,
|
||||
warmup_epochs=args.warmup_epochs,
|
||||
warmup_factor=args.warmup_factor,
|
||||
)
|
||||
loss_function = (
|
||||
egospeak_multihorizon_loss
|
||||
if args.loss == "lstr-mce"
|
||||
else miniroad_multihorizon_loss
|
||||
)
|
||||
|
||||
args.output.parent.mkdir(parents=True, exist_ok=True)
|
||||
history: list[dict[str, Any]] = []
|
||||
best_dev_map = -math.inf if dev_sessions else None
|
||||
best_epoch = args.epochs if not dev_sessions else 0
|
||||
for epoch in range(1, args.epochs + 1):
|
||||
head.train()
|
||||
loss_sum = 0.0
|
||||
example_count = 0
|
||||
progress = tqdm(
|
||||
loader,
|
||||
desc=f"head epoch {epoch}/{args.epochs}",
|
||||
unit="batch",
|
||||
disable=args.no_progress,
|
||||
)
|
||||
for state_batch, target_batch in progress:
|
||||
state_batch = state_batch.to(
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
non_blocking=True,
|
||||
)
|
||||
target_batch = target_batch.to(
|
||||
device=device,
|
||||
dtype=torch.float32,
|
||||
non_blocking=True,
|
||||
)
|
||||
optimizer.zero_grad(set_to_none=True)
|
||||
logits = head(state_batch)
|
||||
loss = loss_function(logits, target_batch)
|
||||
loss.backward()
|
||||
optimizer.step()
|
||||
scheduler.step()
|
||||
batch_count = len(state_batch)
|
||||
loss_sum += float(loss.detach()) * batch_count
|
||||
example_count += batch_count
|
||||
progress.set_postfix(
|
||||
loss=f"{float(loss.detach()):.5f}",
|
||||
lr=f"{scheduler.get_last_lr()[0]:.3g}",
|
||||
)
|
||||
train_loss = loss_sum / example_count
|
||||
record: dict[str, Any] = {
|
||||
"epoch": epoch,
|
||||
"train_loss": train_loss,
|
||||
"learning_rate": scheduler.get_last_lr()[0],
|
||||
}
|
||||
|
||||
if dev_sessions:
|
||||
assert dev_states is not None and dev_targets is not None
|
||||
dev_metrics = _evaluate_valid_horizons(
|
||||
head,
|
||||
dev_states,
|
||||
dev_targets,
|
||||
device,
|
||||
args.batch_size,
|
||||
)
|
||||
dev_map = float(dev_metrics["mean_map"])
|
||||
record["dev_valid_mean_map"] = dev_map
|
||||
print(
|
||||
f"epoch {epoch:02d}: train_loss={train_loss:.6f}, "
|
||||
f"dev_valid_mean_mAP={dev_map * 100:.3f}"
|
||||
)
|
||||
if dev_map > float(best_dev_map):
|
||||
best_dev_map = dev_map
|
||||
best_epoch = epoch
|
||||
torch.save(
|
||||
_checkpoint(
|
||||
head,
|
||||
epoch=epoch,
|
||||
best_epoch=best_epoch,
|
||||
best_dev_map=best_dev_map,
|
||||
train_loss=train_loss,
|
||||
args=args,
|
||||
train_sessions=train_sessions,
|
||||
dev_sessions=dev_sessions,
|
||||
hidden_size=hidden_size,
|
||||
feature_metadata=feature_metadata,
|
||||
history=history + [record],
|
||||
),
|
||||
args.output,
|
||||
)
|
||||
print(f"saved new best checkpoint: {args.output}")
|
||||
else:
|
||||
print(f"epoch {epoch:02d}: train_loss={train_loss:.6f}")
|
||||
history.append(record)
|
||||
|
||||
if not dev_sessions:
|
||||
torch.save(
|
||||
_checkpoint(
|
||||
head,
|
||||
epoch=args.epochs,
|
||||
best_epoch=best_epoch,
|
||||
best_dev_map=None,
|
||||
train_loss=history[-1]["train_loss"],
|
||||
args=args,
|
||||
train_sessions=train_sessions,
|
||||
dev_sessions=dev_sessions,
|
||||
hidden_size=hidden_size,
|
||||
feature_metadata=feature_metadata,
|
||||
history=history,
|
||||
),
|
||||
args.output,
|
||||
)
|
||||
print(f"checkpoint: {args.output}")
|
||||
if dev_sessions:
|
||||
print(
|
||||
f"best epoch: {best_epoch}, "
|
||||
f"dev valid-only mean mAP: {float(best_dev_map) * 100:.3f}"
|
||||
)
|
||||
summary_path = args.output.with_suffix(args.output.suffix + ".json")
|
||||
payload = torch.load(args.output, map_location="cpu", weights_only=False)
|
||||
summary = {key: value for key, value in payload.items() if key != "head"}
|
||||
summary_path.write_text(json.dumps(summary, indent=2) + "\n")
|
||||
print(f"training summary: {summary_path}")
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user