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

Add shared-AV multiwindow Qwen feature extraction

Browse Source
This commit is contained in:
Xin Wang
2026-09-17 03:00:05 +08:00
parent a9f748d6a6
commit 9bc89b556f
8 changed files with 971 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
extract_easycom_multiwindow.py
+179
View File
@@ -0,0 +1,179 @@
#!/usr/bin/env python3
"""Extract EasyCom 2/4/6/8/10/12s context banks with one shared AV pass.
One frozen model is loaded for all selected sessions. The current partial
native block is re-encoded at each 5-Hz cutoff, once rather than once per cap.
Existing direct-media validation and single-window CLI remain unchanged.
"""
from __future__ import annotations
import argparse
import importlib.metadata
import json
import math
import re
import subprocess
import time
from pathlib import Path
from typing import Sequence
import torch
import streaming_turn_model as core
import streaming_turn_model_directly as direct
from streaming_turn_multiwindow import DEFAULT_MEMORY_SECONDS, output_paths, validate_memory_seconds
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("--dataset-root", type=Path, help="EasyCom root containing Main/Video_Compressed")
source.add_argument("--easycom-video-root", type=Path, help="Directory directly containing Session_N")
parser.add_argument("--target-root", type=Path, help="Required for full extraction; target_perframe [T,3]")
parser.add_argument("--sessions", nargs="+", default=[f"Session_{i}" for i in range(1, 13)])
parser.add_argument("--memory-seconds-list", type=float, nargs="+", default=list(DEFAULT_MEMORY_SECONDS))
parser.add_argument("--output-root", type=Path, required=True)
parser.add_argument("--model-id", default=core.MODEL_ID)
parser.add_argument("--backbone-loader", choices=("qwen", "roma"), default="qwen")
parser.add_argument("--dtype", choices=("auto", "bfloat16", "float16", "float32"), default="bfloat16")
parser.add_argument("--device", default="cuda:0")
parser.add_argument("--flash-attention-2", action="store_true")
parser.add_argument("--end", type=float, help="Bounded zero-origin smoke only; does not attach full targets")
parser.add_argument("--readout-last-n-layers", type=int, default=4)
parser.add_argument("--readout-layer-fusion", choices=("mean", "stack"), default="stack")
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("--decoder", choices=("decord", "torchcodec", "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=2)
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("--media-only", action="store_true")
parser.add_argument("--overwrite", action="store_true")
return parser.parse_args(argv)
def session_commands(args: argparse.Namespace) -> list[tuple[str, list[str]]]:
"""Preflight all paths/settings before loading the backbone."""
caps = validate_memory_seconds(args.memory_seconds_list, 2.0)
core.StreamConfig(
memory_seconds=max(caps), video_fps=args.video_fps,
video_min_pixels=args.video_min_pixels, video_max_pixels=args.video_max_pixels,
readout_last_n_layers=args.readout_last_n_layers, readout_layer_fusion=args.readout_layer_fusion,
).validate()
if len(set(args.sessions)) != len(args.sessions) or not args.sessions:
raise ValueError("sessions must be nonempty and unique")
if args.end is None and args.target_root is None:
raise ValueError("full extraction requires --target-root; use --end only for smoke")
if args.end is not None:
if not math.isfinite(args.end) or args.end <= 0:
raise ValueError("smoke end must be finite and positive")
core._require_integer_ratio(args.end, 0.2, "duration / step")
if args.verify_every < 0 or args.verify_continuous_every < 0:
raise ValueError("verification intervals must be non-negative")
if args.decoder_threads <= 0 or args.max_open_readers <= 0 or args.frame_cache_size < 0:
raise ValueError("invalid reader/cache limits")
video_root = args.easycom_video_root or args.dataset_root / "Main" / "Video_Compressed"
commands = []
for session in args.sessions:
if re.fullmatch(r"Session_(?:[1-9]|1[0-2])", session) is None:
raise ValueError(f"invalid EasyCom session: {session}")
session_dir = video_root / session
if not session_dir.is_dir():
raise FileNotFoundError(session_dir)
direct._easycom_directory_specs(session_dir)
output = args.output_root / f"{session}.pt"
if not args.overwrite and not args.media_only:
for path in output_paths(output, caps).values():
if path.exists():
raise FileExistsError(f"{path}; use --overwrite explicitly or a fresh output root")
command = ["--easycom-session-dir", str(session_dir), "--output", str(output),
"--mode", "open-stream-replay", "--memory-seconds-list", *map(str, caps)]
for name in (
"model_id", "backbone_loader", "dtype", "device", "readout_last_n_layers",
"readout_layer_fusion", "video_fps", "video_min_pixels", "video_max_pixels",
"decoder", "decoder_gpu_id", "decoder_threads", "max_open_readers", "frame_cache_size",
"verify_every", "verify_continuous_every",
):
command.extend(["--" + name.replace("_", "-"), str(getattr(args, name))])
for name in ("flash_attention_2", "no_progress", "media_only", "overwrite"):
if getattr(args, name):
command.append("--" + name.replace("_", "-"))
if args.end is not None:
command.extend(["--end", str(args.end)])
else:
target = args.target_root / f"{session}.npy"
core.target_num_steps(target) # Missing/invalid targets fail before model loading.
command.extend(["--target", str(target)])
commands.append((session, command))
return commands
def main(argv: Sequence[str] | None = None) -> None:
args = parse_args(argv)
commands = session_commands(args)
if args.media_only:
for _, command in commands:
direct.main(command)
return
if args.end is not None:
print("SMOKE ONLY: partial banks are not a full-session training/evaluation dataset.", flush=True)
core._configure_qwen_logging()
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, backbone_loader=args.backbone_loader,
)
model_loading_seconds = time.perf_counter() - load_started
repository = Path(__file__).resolve().parent
revision = subprocess.run(["git", "rev-parse", "HEAD"], cwd=repository, capture_output=True, text=True, check=False)
dirty = subprocess.run(["git", "status", "--porcelain"], cwd=repository, capture_output=True, text=True, check=False)
summary = {
"format": "easycom-multiwindow-extraction-v1", "arguments": {
key: str(value) if isinstance(value, Path) else value for key, value in vars(args).items()
},
"split": {"train": [f"Session_{i}" for i in range(4, 13)],
"test": ["Session_1", "Session_2", "Session_3"]},
"backbone": "frozen", "partial_extraction": args.end is not None,
"code_revision": revision.stdout.strip() if revision.returncode == 0 else None,
"worktree_dirty": bool(dirty.stdout.strip()) if dirty.returncode == 0 else None,
"versions": {name: importlib.metadata.version(name) for name in ("torch", "transformers", "qwen-omni-utils")},
"model_loading_seconds": model_loading_seconds,
"session_wall_scope": "direct_media_index_decode_extract_save_excluding_model_load_and_summary_scan",
"runtime_scope": "joint_multi_bank_with_shared_AV_not_single_window_online_runtime",
"sessions": {},
}
for session, command in commands:
started = time.perf_counter()
direct.main(command, model=model, processor=processor)
elapsed = time.perf_counter() - started
records = {}
for cap, path in output_paths(args.output_root / f"{session}.pt", sorted(args.memory_seconds_list)).items():
payload = torch.load(path, map_location="cpu", weights_only=False)
checks = [c for c in payload["verification"] if c.get("kind") == "cache_replay"]
spans = payload["extraction_diagnostics"]["actual_context_seconds"]
records[f"{cap:g}s"] = {
"path": str(path), "shape": list(payload["states"].shape),
"num_steps": payload["num_steps"], "max_actual_context_seconds": spans.max().item(),
"replay_checks": len(checks),
"worst_replay_max_abs": max((c["max_abs"] for c in checks), default=None),
"lowest_replay_cosine": min((c["cosine"] for c in checks), default=None),
"production_processor_calls_shared": payload["num_steps"],
}
del payload
summary["sessions"][session] = {"joint_wall_seconds": elapsed, "banks": records}
args.output_root.mkdir(parents=True, exist_ok=True)
summary_path = args.output_root / "extraction_summary.json"
temporary = summary_path.with_suffix(".json.tmp")
temporary.write_text(json.dumps(summary, indent=2) + "\n", encoding="utf-8")
temporary.replace(summary_path)
print(f"Summary: {args.output_root / 'extraction_summary.json'}")
if __name__ == "__main__":
main()