wx44wx/icassp2027-experiment
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Add reproducible EasyCom causal head pipeline

Browse Source
This commit is contained in:
Xin Wang
2026-08-16 22:19:43 +08:00
parent bde72487bb
commit 080c45f049
8 changed files with 1451 additions and 470 deletions
Download Patch File Download Diff File Expand all files Collapse all files

165
README.md
View File

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

299
easycom_head.py Normal file
View 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
View 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()

76
prepare_easycom_session.py
View File

@@ -42,10 +42,22 @@ class ClipInfo:
def _run_json(command: list[str]) -> dict:
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,

2
requirements.txt
View File

@@ -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
View 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

533
streaming_turn_model.py
View File

@@ -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:
@@ -167,19 +209,22 @@ class FixedGridAV:
)
def frames(self, start: float, end: float) -> list[Any]:
"""Return fixed-grid frames whose timestamps lie in [start, end]."""
"""Return fixed-grid frames whose timestamps lie in [start, end)."""
import numpy as np
from PIL import Image
first = max(0, int(math.ceil(start * self.fps - 1e-7)))
last = min(
self.num_frames - 1,
int(math.floor(end * self.fps + 1e-7)),
last_exclusive = min(
self.num_frames,
int(math.ceil(end * self.fps - 1e-7)),
)
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()
@@ -198,76 +243,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)
@@ -437,16 +412,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)
@@ -662,18 +644,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:
@@ -823,13 +797,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,
@@ -837,13 +820,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,
@@ -876,32 +870,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:
@@ -982,11 +985,23 @@ 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 = []
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)
context_start = max(
start, current_time - self.config.memory_seconds
)
state = self.state_for_intervals(
video_path=video_path,
audio_path=audio_path,
@@ -994,6 +1009,10 @@ class CausalPrefixExtractor:
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),
@@ -1001,202 +1020,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."""
@@ -1271,7 +1094,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,
@@ -1280,8 +1109,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,
)
@@ -1292,7 +1121,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:
@@ -1300,6 +1129,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):
@@ -1312,7 +1142,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,
@@ -1322,9 +1156,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,
@@ -1337,6 +1179,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
@@ -1373,6 +1216,7 @@ class OpenStreamReplayExtractor:
_,
payload_base,
_,
_,
) = self._cache_from_intervals(
video_path=video_path,
audio_path=audio_path,
@@ -1394,7 +1238,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,
@@ -1404,8 +1252,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())
@@ -1454,6 +1302,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(
@@ -1466,7 +1319,18 @@ class OpenStreamReplayExtractor:
use_audio_in_video: bool,
verify_every: int = 0,
verify_continuous_every: int = 0,
show_progress: bool = True,
) -> ExtractionResult:
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,
@@ -1476,6 +1340,7 @@ class OpenStreamReplayExtractor:
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),
@@ -1485,7 +1350,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
@@ -1637,7 +1502,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)
@@ -1661,6 +1526,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",
@@ -1674,6 +1544,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")
@@ -1683,7 +1554,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:
@@ -1758,19 +1628,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(
@@ -1784,6 +1642,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] = {
@@ -1792,11 +1651,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"
)
),
"config": asdict(config),
"start": args.start,
@@ -1808,7 +1663,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"
),

476
train_turn_head.py Normal file
View 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()