Add multi-GPU Ego4D local feature extraction without symlinks

run_ego4d_local_baselines.py

@@ -13,7 +13,7 @@ import shlex
 import subprocess
 import sys
 
-from easycom_head import load_ego4d_manifest_split, load_session_data
+from easycom_head import load_ego4d_manifest_split, load_session_data, resolve_feature_path
 
 
 VARIANTS = ("baseline", "audio", "visual", "av")
@@ -24,7 +24,8 @@ def parse_args(argv=None):
     p = argparse.ArgumentParser(description=__doc__)
     p.add_argument("--data-root", type=Path, default=Path("/data/wangx/Datasets/ego4d"))
     p.add_argument("--manifest", type=Path)
-    p.add_argument("--feature-root", required=True, type=Path, help="Your EXISTING global Omni features")
+    p.add_argument("--feature-root", required=True, type=Path,
+                   help="Existing global Omni features: flat or Ego4D train/test parent")
     p.add_argument("--target-root", type=Path)
     p.add_argument("--local-feature-root", type=Path)
     p.add_argument("--checkpoint-root", type=Path, default=Path("checkpoints/ego4d_local_b1"))
@@ -67,11 +68,19 @@ def preflight(args) -> str:
     use_local = any(variant != "baseline" for variant in args.variants)
     wanted = set(args.variants) - {"baseline"}
     local_modalities = next(iter(wanted)) if len(wanted) == 1 else "av"
-    roots = [(args.feature_root, ".pt"), (args.target_root, ".npy")]
+    roots = [args.feature_root]
     if use_local:
-        roots.append((args.local_feature_root, ".pt"))
-    missing = [str(root / f"{name}{suffix}") for root, suffix in roots for name in names
-               if not (root / f"{name}{suffix}").is_file()]
+        roots.append(args.local_feature_root)
+    missing = []
+    for name in names:
+        target = args.target_root / f"{name}.npy"
+        if not target.is_file():
+            missing.append(str(target))
+        for root in roots:
+            try:
+                resolve_feature_path(root, name, dataset="ego4d")
+            except FileNotFoundError as error:
+                missing.append(str(error))
     if missing:
         raise FileNotFoundError(f"Missing {len(missing)} input files (first 20):\n" + "\n".join(missing[:20]))
     for variant in args.variants: