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

Add extract_ego4d_batch.py for efficient EgoSpeak clip extraction. This script loads the Qwen model once and processes selected manifest rows sequentially, ensuring each clip is handled independently. Update README with usage instructions and details on the extraction process, including handling of completed outputs and memory management.

Browse Source
This commit is contained in:
Xin Wang
2026-08-18 22:22:37 +08:00
parent 609b2781d4
commit 5fa61eed23
3 changed files with 514 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
README.md
View File

@@ -685,6 +685,38 @@ conversion can differ slightly; validate a short Decord/PyNvVideoCodec pair
before extracting every session, and do not mix decoder backends across
train/test splits.
### Persistent-model Ego4D batch extraction
`extract_ego4d_batch.py` loads Qwen and its processor once, then extracts the
selected EgoSpeak manifest rows sequentially. Each clip receives a new media
reader, Thinker KV cache, TMRoPE origin, and rolling-memory state; only the
immutable model weights and processor are reused. Manifest indices are
one-based data-row indices and do not count the TSV header.
For example, process manifest samples 81 through the end on visible GPU 0:
```bash
CUDA_VISIBLE_DEVICES=0 \
OMP_NUM_THREADS=4 \
MKL_NUM_THREADS=4 \
OPENBLAS_NUM_THREADS=4 \
PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True \
python extract_ego4d_batch.py \
--data-root /data1/yfzx/wangx/Datasets/ego4d \
--model-id /data1/yfzx/wangx/Models/Qwen2.5-Omni-3B \
--start-index 81 \
--memory-seconds 20 \
--video-fps 2 \
--decoder decord \
--dtype bfloat16 \
--flash-attention-2 \
--device cuda:0
```
Completed non-empty `.pt` files are skipped. Each new output is first written
to a process-specific `.partial` path and atomically renamed after extraction,
so rerunning the same range resumes safely.
For CPU Decord, the direct reader intentionally indexes each clip and releases
that temporary reader, then keeps an LRU of at most two live readers during
extraction. Each reader defaults to one decode thread. These limits avoid one

454
extract_ego4d_batch.py Normal file
View File

@@ -0,0 +1,454 @@
#!/usr/bin/env python3
"""Extract EgoSpeak's Ego4D clips while loading Qwen only once.
Manifest sample indices are one-based and refer to data rows, excluding the
header. For example, ``--start-index 31 --end-index 60`` selects exactly the
same clips as ``sed -n '32,61p' ego4d_egospeak.tsv``.
Only immutable model weights and the processor are reused. Every clip invokes
the direct extractor independently, which creates a new media reader, rolling
KV cache, TMRoPE origin, and readout sequence. Consequently no temporal state
is carried between dataset samples.
"""
from __future__ import annotations
import argparse
import csv
import gc
import math
import os
import sys
import time
import traceback
from dataclasses import dataclass
from pathlib import Path
from typing import Sequence
import torch
import streaming_turn_model as core
import streaming_turn_model_directly as direct
TARGET_HZ = 5
DEFAULT_OUTPUT_NAME = "open_stream_2fps_20s_decord"
EXPECTED_MANIFEST_COLUMNS = (
"split",
"uid",
"steps",
"duration",
"video_frames_2fps",
)
@dataclass(frozen=True)
class ClipRecord:
index: int
split: str
uid: str
steps: int
duration: float
video_frames_2fps: int
def parse_args(argv: Sequence[str] | None = None) -> argparse.Namespace:
parser = argparse.ArgumentParser(description=__doc__)
parser.add_argument(
"--data-root",
type=Path,
default=Path("/data/wangx/Datasets/ego4d"),
help="Ego4D experiment root (default: %(default)s)",
)
parser.add_argument(
"--manifest",
type=Path,
help="Defaults to DATA_ROOT/ego4d_egospeak.tsv",
)
parser.add_argument(
"--video-root",
type=Path,
help="Defaults to DATA_ROOT/v1/clips",
)
parser.add_argument(
"--target-root",
type=Path,
help="Defaults to DATA_ROOT/ego4d_target_frame",
)
parser.add_argument(
"--output-root",
type=Path,
help=(
"Defaults to DATA_ROOT/omni_features/"
f"{DEFAULT_OUTPUT_NAME}"
),
)
parser.add_argument(
"--model-id",
default="/data/wangx/Models/Qwen2.5-Omni-3B",
)
parser.add_argument(
"--start-index",
type=int,
default=1,
help="First one-based manifest data row to process (default: 1)",
)
parser.add_argument(
"--end-index",
type=int,
help="Last one-based manifest data row to process, inclusive",
)
parser.add_argument(
"--split",
choices=("all", "train", "test"),
default="all",
)
parser.add_argument(
"--overwrite",
action="store_true",
help="Replace completed outputs instead of resuming around them",
)
parser.add_argument(
"--continue-on-error",
action="store_true",
help="Record a failed clip and continue with the next selected clip",
)
parser.add_argument("--memory-seconds", type=float, default=20.0)
parser.add_argument("--step-seconds", type=float, default=0.2)
parser.add_argument("--native-block-seconds", type=float, default=2.0)
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=64)
parser.add_argument(
"--decoder",
choices=("decord", "pynvcodec"),
default="decord",
)
parser.add_argument("--decoder-gpu-id", type=int, default=0)
parser.add_argument("--decoder-threads", type=int, default=1)
parser.add_argument("--max-open-readers", type=int, default=1)
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(
"--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(
"--empty-cache-every",
type=int,
default=0,
help="Call torch.cuda.empty_cache every N processed clips; 0 disables",
)
args = parser.parse_args(argv)
if args.start_index <= 0:
parser.error("--start-index must be positive")
if args.end_index is not None and args.end_index < args.start_index:
parser.error("--end-index must be >= --start-index")
if args.empty_cache_every < 0:
parser.error("--empty-cache-every must be >= 0")
if args.step_seconds <= 0:
parser.error("--step-seconds must be positive")
return args
def resolve_paths(args: argparse.Namespace) -> None:
args.data_root = args.data_root.expanduser().resolve()
args.manifest = (
args.manifest or args.data_root / "ego4d_egospeak.tsv"
).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_features" / DEFAULT_OUTPUT_NAME
).expanduser().resolve()
def load_manifest(path: Path) -> list[ClipRecord]:
if not path.is_file():
raise FileNotFoundError(f"Missing EgoSpeak manifest: {path}")
records: list[ClipRecord] = []
seen: set[str] = set()
with path.open(newline="") as file:
reader = csv.DictReader(file, delimiter="\t")
if reader.fieldnames != list(EXPECTED_MANIFEST_COLUMNS):
raise RuntimeError(
f"Unexpected manifest columns in {path}: {reader.fieldnames}"
)
for index, row in enumerate(reader, start=1):
split = row["split"].strip()
uid = row["uid"].strip()
if split not in {"train", "test"}:
raise RuntimeError(f"Invalid split at manifest row {index}: {split}")
if not uid or uid in seen:
raise RuntimeError(
f"Missing or duplicate UID at manifest row {index}: {uid!r}"
)
seen.add(uid)
steps = int(row["steps"])
duration = float(row["duration"])
video_frames = int(row["video_frames_2fps"])
if steps <= 0:
raise RuntimeError(f"Non-positive steps for {uid}: {steps}")
expected_duration = steps / TARGET_HZ
if not math.isclose(
duration, expected_duration, rel_tol=0.0, abs_tol=1e-6
):
raise RuntimeError(
f"Manifest duration mismatch for {uid}: {duration} vs "
f"{expected_duration} from {steps} target steps"
)
records.append(
ClipRecord(
index=index,
split=split,
uid=uid,
steps=steps,
duration=duration,
video_frames_2fps=video_frames,
)
)
if not records:
raise RuntimeError(f"No clips in {path}")
return records
def select_records(
records: list[ClipRecord], args: argparse.Namespace
) -> list[ClipRecord]:
end_index = args.end_index if args.end_index is not None else len(records)
selected = [
record
for record in records
if args.start_index <= record.index <= end_index
and (args.split == "all" or record.split == args.split)
]
if not selected:
raise RuntimeError(
f"No manifest rows selected by {args.start_index=} "
f"{args.end_index=} {args.split=}"
)
return selected
def preflight(selected: list[ClipRecord], args: argparse.Namespace) -> None:
if not args.video_root.is_dir():
raise FileNotFoundError(f"Missing video root: {args.video_root}")
if not args.target_root.is_dir():
raise FileNotFoundError(f"Missing target root: {args.target_root}")
model_path = Path(args.model_id).expanduser()
if model_path.is_absolute() and not model_path.is_dir():
raise FileNotFoundError(f"Missing local model: {model_path}")
problems: list[str] = []
for record in selected:
video = args.video_root / f"{record.uid}.mp4"
target = args.target_root / f"{record.uid}.npy"
if not video.is_file() or video.stat().st_size == 0:
problems.append(f"missing video: {video}")
if not target.is_file() or target.stat().st_size == 0:
problems.append(f"missing target: {target}")
continue
try:
target_steps = core.target_num_steps(target)
except Exception as error: # report all bad inputs before model loading
problems.append(f"invalid target {target}: {error}")
continue
if target_steps != record.steps:
problems.append(
f"target length mismatch {record.uid}: manifest={record.steps}, "
f"npy={target_steps}"
)
if problems:
preview = "\n".join(f"- {problem}" for problem in problems[:20])
suffix = "" if len(problems) <= 20 else f"\n... and {len(problems)-20} more"
raise RuntimeError(f"Ego4D preflight failed:\n{preview}{suffix}")
def clip_argv(
record: ClipRecord,
args: argparse.Namespace,
output: Path,
) -> list[str]:
values = [
"--model-id",
str(args.model_id),
"--video",
str(args.video_root / f"{record.uid}.mp4"),
"--target",
str(args.target_root / f"{record.uid}.npy"),
"--mode",
"open-stream-replay",
"--step-seconds",
str(args.step_seconds),
"--native-block-seconds",
str(args.native_block_seconds),
"--memory-seconds",
str(args.memory_seconds),
"--video-fps",
str(args.video_fps),
"--video-min-pixels",
str(args.video_min_pixels),
"--video-max-pixels",
str(args.video_max_pixels),
"--frame-cache-size",
str(args.frame_cache_size),
"--decoder",
args.decoder,
"--decoder-gpu-id",
str(args.decoder_gpu_id),
"--decoder-threads",
str(args.decoder_threads),
"--max-open-readers",
str(args.max_open_readers),
"--verify-every",
str(args.verify_every),
"--verify-continuous-every",
str(args.verify_continuous_every),
"--dtype",
args.dtype,
"--device",
args.device,
"--output",
str(output),
]
if args.flash_attention_2:
values.append("--flash-attention-2")
if args.no_progress:
values.append("--no-progress")
return values
def format_eta(seconds: float) -> str:
if not math.isfinite(seconds) or seconds < 0:
return "unknown"
hours, remainder = divmod(int(seconds), 3600)
minutes, secs = divmod(remainder, 60)
return f"{hours:02d}:{minutes:02d}:{secs:02d}"
def main(argv: Sequence[str] | None = None) -> None:
args = parse_args(argv)
resolve_paths(args)
core._configure_qwen_logging()
records = load_manifest(args.manifest)
selected = select_records(records, args)
preflight(selected, args)
(args.output_root / "train").mkdir(parents=True, exist_ok=True)
(args.output_root / "test").mkdir(parents=True, exist_ok=True)
pending = [
record
for record in selected
if args.overwrite
or not (args.output_root / record.split / f"{record.uid}.pt").is_file()
or (args.output_root / record.split / f"{record.uid}.pt").stat().st_size == 0
]
skipped = len(selected) - len(pending)
print(
f"Selected manifest samples {selected[0].index}-{selected[-1].index}: "
f"{len(selected)} clip(s), pending={len(pending)}, skipped={skipped}",
flush=True,
)
print(f"Output root: {args.output_root}", flush=True)
if not pending:
print("All selected clips already have non-empty outputs.")
return
print(f"Loading Qwen once from {args.model_id} ...", flush=True)
load_started = time.perf_counter()
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)
print(
f"Model loaded once in {time.perf_counter() - load_started:.1f}s; "
"per-clip KV/media state will reset.",
flush=True,
)
run_started = time.perf_counter()
completed = 0
failures: list[tuple[ClipRecord, str]] = []
for pending_index, record in enumerate(pending, start=1):
output = args.output_root / record.split / f"{record.uid}.pt"
temporary = output.with_name(f"{output.name}.partial.{os.getpid()}")
print(
f"\n[{pending_index}/{len(pending)} | manifest {record.index}] "
f"{record.split}/{record.uid}: {record.steps} states, "
f"{record.duration:.1f}s",
flush=True,
)
clip_started = time.perf_counter()
try:
direct.main(
clip_argv(record, args, temporary),
model=model,
processor=processor,
)
os.replace(temporary, output)
completed += 1
elapsed = time.perf_counter() - clip_started
average = (time.perf_counter() - run_started) / completed
remaining = len(pending) - pending_index
print(
f"Completed {output} in {elapsed:.1f}s; "
f"ETA {format_eta(average * remaining)}",
flush=True,
)
except KeyboardInterrupt:
print(f"Interrupted; partial path: {temporary}", file=sys.stderr)
raise
except Exception as error:
failures.append((record, str(error)))
print(
f"FAILED manifest sample {record.index} ({record.uid}): {error}\n"
f"Partial path, if created: {temporary}",
file=sys.stderr,
flush=True,
)
traceback.print_exc()
if not args.continue_on_error:
raise
finally:
# The model and processor stay referenced above. Collection here
# releases the completed clip's reader, PIL frames, branch caches,
# and Python cache wrappers before opening the next sample.
gc.collect()
processed = completed + len(failures)
if (
args.empty_cache_every > 0
and processed > 0
and processed % args.empty_cache_every == 0
and torch.cuda.is_available()
):
torch.cuda.empty_cache()
total = time.perf_counter() - run_started
print(
f"\nBatch finished in {format_eta(total)}: completed={completed}, "
f"skipped={skipped}, failed={len(failures)}",
flush=True,
)
if failures:
for record, message in failures:
print(f"- manifest {record.index} {record.uid}: {message}", file=sys.stderr)
raise SystemExit(1)
if __name__ == "__main__":
main()

39
streaming_turn_model_directly.py
View File

@@ -59,7 +59,7 @@ import subprocess
from collections import OrderedDict
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any
from typing import Any, Sequence
import numpy as np
import torch
@@ -778,7 +778,7 @@ def _parse_crop(value: str | None) -> tuple[int, int, int, int] | None:
return values
def parse_args() -> argparse.Namespace:
def parse_args(argv: Sequence[str] | None = None) -> 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")
@@ -873,7 +873,7 @@ def parse_args() -> argparse.Namespace:
action="store_true",
help="Validate/index direct media and exit before loading Qwen",
)
return parser.parse_args()
return parser.parse_args(argv)
def _resolve_duration_and_target(
@@ -941,9 +941,23 @@ def _print_media_summary(reader: CausalSourceAV, start: float, end: float) -> No
)
def main() -> None:
args = parse_args()
def main(
argv: Sequence[str] | None = None,
*,
model: Any | None = None,
processor: Any | None = None,
) -> None:
"""Run one direct extraction, optionally reusing loaded Qwen objects.
``model`` and ``processor`` must be supplied together. Reusing them is
safe across independent clips because this function constructs a new media
reader, extractor, and temporal cache on every call.
"""
args = parse_args(argv)
core._configure_qwen_logging()
if (model is None) != (processor is None):
raise ValueError("model and processor must be supplied together")
if args.verify_every < 0 or args.verify_continuous_every < 0:
raise SystemExit("verification intervals must be >= 0")
if args.frame_cache_size < 0:
@@ -1055,12 +1069,15 @@ def main() -> None:
)
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,
)
if model is None:
model, processor = core.load_qwen_omni(
args.model_id,
dtype=args.dtype,
device_map={"": args.device},
use_flash_attention_2=args.flash_attention_2,
)
else:
model.eval()
model.requires_grad_(False)
media_path_for_api = reader.clips[0].path