Add multi-GPU Ego4D local feature extraction without symlinks

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

View File

@@ -18,7 +18,7 @@ from unittest.mock import patch
import numpy as np
import torch
from easycom_head import load_ego4d_manifest_split
from easycom_head import load_ego4d_manifest_split, load_session_data, resolve_feature_path
import extract_turn_local_features as extract
import run_ego4d_local_baselines as runner
from test_local_turn_features import fixture
@@ -73,6 +73,45 @@ class Ego4DLocalWorkflowTests(unittest.TestCase):
with self.assertRaisesRegex(ValueError, "Duplicate"):
load_ego4d_manifest_split(self.manifest, "train")
def split_feature_layout(self, folder="features"):
for name, split in zip(self.names, self.splits):
destination = self.root / folder / split
destination.mkdir(exist_ok=True)
(self.root / folder / f"{name}.pt").rename(destination / f"{name}.pt")
def test_native_split_paths_and_ambiguous_copies(self):
self.split_feature_layout()
self.split_feature_layout("local")
for name, split in zip(self.names, self.splits):
path = resolve_feature_path(self.root / "features", name, dataset="ego4d")
self.assertEqual(path, self.root / "features" / split / f"{name}.pt")
self.assertFalse(path.is_symlink())
data = load_session_data(self.root / "features", self.root / "targets", name,
dataset="ego4d", expected_mode="open-stream-replay",
local_feature_root=self.root / "local")
self.assertEqual(len(data.states), 20)
duplicate = self.root / "features" / f"{self.names[0]}.pt"
duplicate.write_bytes(b"must-not-be-silently-chosen")
with self.assertRaisesRegex(ValueError, "Ambiguous feature copies"):
resolve_feature_path(self.root / "features", self.names[0], dataset="ego4d")
def test_shards_are_disjoint_and_cover_filtered_selection(self):
selected = []
for rank in range(8):
args = extract.parse_args(self.extraction_args(["--num-shards", "8", "--shard-index", str(rank)]))
shard = extract.select_samples(args)
self.assertEqual(shard, self.names[rank::8])
selected.extend(shard)
self.assertCountEqual(selected, self.names)
for rank in range(2):
args = extract.parse_args(self.extraction_args(["--split", "train", "--num-shards", "2",
"--shard-index", str(rank)]))
self.assertEqual(extract.select_samples(args), [self.names[rank * 2]])
def test_split_input_directory_is_protected_from_output_overwrite(self):
with contextlib.redirect_stderr(io.StringIO()), self.assertRaises(SystemExit):
extract.parse_args(self.extraction_args(["--output-root", str(self.root / "features/train")]))
def test_invalid_cli_combinations(self):
for extra in (["--samples", self.names[0]], ["--dataset", "easycom"], ["--start-index", "0"],
["--limit", "0"], ["--max-steps", "20", "--reuse-existing"],
@@ -94,6 +133,7 @@ class Ego4DLocalWorkflowTests(unittest.TestCase):
extract.main(self.extraction_args(["--dry-run"]))
def test_one_encoder_per_batch_and_validated_reuse(self):
self.split_feature_layout()
fake_encoder = SimpleNamespace(transformers_version="4.57.6", visual_dim=6, audio_dim=5,
model_config_sha256="same-checkpoint")
_, side, _ = fixture(dataset="ego4d")
@@ -174,6 +214,7 @@ class Ego4DLocalWorkflowTests(unittest.TestCase):
self.assertFalse((self.root / "checkpoints").exists())
def test_full_cpu_family_round_trip_and_summary(self):
self.split_feature_layout()
env = dict(os.environ, OMP_NUM_THREADS="1", MKL_NUM_THREADS="1", OPENBLAS_NUM_THREADS="1")
completed = subprocess.run([sys.executable, str(runner.REPO / "run_ego4d_local_baselines.py"),
*self.runner_args()], env=env, capture_output=True, text=True, timeout=90)