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

Add multi-GPU Ego4D local feature extraction without symlinks

Browse Source
This commit is contained in:
Xin Wang
2026-08-28 01:01:25 +08:00
parent 9bd204b468
commit 5fd285e369
8 changed files with 552 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

28
extract_turn_local_features.py
View File

@@ -19,7 +19,7 @@ import numpy as np
import torch
from tqdm.auto import tqdm
from easycom_head import load_ego4d_manifest_split, normalize_sample_name
from easycom_head import load_ego4d_manifest_split, normalize_sample_name, resolve_feature_path
from local_turn_features import (
LOCAL_FORMAT, causal_frame_ticks, crop_box, easycom_boxes,
load_local_features, pooled_recent_visual, reference_alignment,
@@ -406,6 +406,9 @@ def select_samples(args: argparse.Namespace) -> list[str]:
raise ValueError("No samples selected")
if len(set(names)) != len(names):
raise ValueError("Duplicate samples")
num_shards = getattr(args, "num_shards", 1)
shard_index = getattr(args, "shard_index", 0)
names = names[shard_index::num_shards]
return names
@@ -445,7 +448,8 @@ def parse_args(argv=None):
p = argparse.ArgumentParser(description=__doc__)
p.add_argument("--dataset", choices=("easycom", "ego4d"), default="easycom")
p.add_argument("--data-root", required=True, type=Path)
p.add_argument("--feature-root", required=True, type=Path)
p.add_argument("--feature-root", required=True, type=Path,
help="Flat .pt directory or Ego4D parent containing train/ and test/")
p.add_argument("--target-root", required=True, type=Path)
p.add_argument("--output-root", required=True, type=Path)
selection = p.add_mutually_exclusive_group(required=True)
@@ -455,6 +459,8 @@ def parse_args(argv=None):
p.add_argument("--start-index", type=int, default=1, help="One-based index AFTER split filtering")
p.add_argument("--end-index", type=int, help="Inclusive index AFTER split filtering")
p.add_argument("--limit", type=int, help="First N clips in the selected manifest range")
p.add_argument("--num-shards", type=int, default=1, help="Disjoint round-robin shards after all sample selection")
p.add_argument("--shard-index", type=int, default=0, help="Zero-based shard index")
p.add_argument("--ego4d-annotations", type=Path, nargs="+")
p.add_argument("--model-id", default="Qwen/Qwen2.5-Omni-3B")
p.add_argument("--encoder-id", help="Stable checkpoint identity for relocated copies; default: model-id")
@@ -490,6 +496,8 @@ def parse_args(argv=None):
p.add_argument("--dry-run", action="store_true", help="Validate references/targets and list work; no model, media decode or writes")
p.add_argument("--no-progress", action="store_true")
args = p.parse_args(argv)
if args.num_shards <= 0 or not 0 <= args.shard_index < args.num_shards:
p.error("Require --num-shards > 0 and 0 <= --shard-index < --num-shards")
if args.manifest is not None and args.dataset != "ego4d":
p.error("--manifest is only supported with --dataset ego4d")
if args.manifest is None and (args.split != "all" or args.start_index != 1
@@ -527,7 +535,10 @@ def parse_args(argv=None):
p.error("max-pixels must be >=3136; image-batch-size must be positive")
if args.max_steps is not None and args.max_steps <= 0:
p.error("--max-steps must be positive")
if args.output_root.resolve() == args.feature_root.resolve():
protected_roots = [args.feature_root]
if args.dataset == "ego4d":
protected_roots += [args.feature_root / split for split in ("train", "test")]
if args.output_root.resolve() in [root.resolve() for root in protected_roots]:
p.error("Local outputs must not overwrite original Omni features")
return args
@@ -541,6 +552,11 @@ def main(argv=None) -> None:
# Fail before indexing/decoding source media, not after expensive I/O.
version = validate_local_transformers_version(transformers.__version__)
names = select_samples(args)
if not names:
print(f"Shard {args.shard_index}/{args.num_shards}: no assigned clips")
return
print(f"Shard {args.shard_index}/{args.num_shards}: {len(names)} assigned clips", flush=True)
reference_paths = {name: resolve_feature_path(args.feature_root, name, dataset=args.dataset) for name in names}
paths = []
if args.dataset == "ego4d":
paths = args.ego4d_annotations or [args.data_root / "v1" / "annotations" / f"av_{split}.json"
@@ -549,7 +565,7 @@ def main(argv=None) -> None:
if not path.is_file():
raise FileNotFoundError(path)
for name in names:
required = [args.feature_root / f"{name}.pt", args.target_root / f"{name}.npy"]
required = [reference_paths[name], args.target_root / f"{name}.npy"]
if args.dataset == "ego4d":
required.append(args.data_root / "v1" / "clips" / f"{name}.mp4")
for path in required:
@@ -561,7 +577,7 @@ def main(argv=None) -> None:
pending, contract = [], None
# Stream through references on CPU; never keep all global states in memory.
for name in tqdm(names, desc="local preflight", unit="clip", disable=args.no_progress or args.dry_run):
reference = torch.load(args.feature_root / f"{name}.pt", map_location="cpu", weights_only=False)
reference = torch.load(reference_paths[name], map_location="cpu", weights_only=False)
times, _ = reference_alignment(reference)
target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")
if target.shape != (len(times), 3):
@@ -586,7 +602,7 @@ def main(argv=None) -> None:
encoder = None
for position, name in enumerate(pending, 1):
print(f"[{position}/{len(pending)}] {name}", flush=True)
path = args.feature_root / f"{name}.pt"
path = reference_paths[name]
reference = torch.load(path, map_location="cpu", weights_only=False)
times, indices = reference_alignment(reference)
target = np.load(args.target_root / f"{name}.npy", allow_pickle=False, mmap_mode="r")