"""Run Ego4D G / G+A / G+V / G+A+V: four epochs, no dev, manifest splits. All jobs run in the foreground. The full selected family is preflighted before any training, including test feature contracts. No test metrics select epochs. """ from __future__ import annotations import argparse import os from pathlib import Path import shlex import subprocess import sys from easycom_head import load_ego4d_manifest_split, load_session_data, resolve_feature_path VARIANTS = ("baseline", "audio", "visual", "av") REPO = Path(__file__).resolve().parent 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="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")) p.add_argument("--result-root", type=Path, default=Path("results/ego4d_local_b1")) p.add_argument("--variants", choices=VARIANTS, nargs="+", default=list(VARIANTS)) p.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2]) p.add_argument("--layer-fusion", choices=("auto", "single", "concat-mlp", "fixed-mean", "scalar-mixture"), default="auto", help="auto: concat-mlp for stacked states, single for [T,D]") p.add_argument("--local-projection-dim", type=int, default=256) p.add_argument("--expected-mode", choices=("open-stream-replay", "causal-prefix"), default="open-stream-replay") p.add_argument("--device", default="cuda:0", help="Visible device; select the physical GPU with CUDA_VISIBLE_DEVICES") p.add_argument("--no-progress", action="store_true") p.add_argument("--dry-run", action="store_true", help="Read/validate all inputs and print commands; no training or writes") args = p.parse_args(argv) args.manifest = args.manifest or args.data_root / "ego4d_egospeak.tsv" args.target_root = args.target_root or args.data_root / "ego4d_target_frame" args.local_feature_root = args.local_feature_root or args.data_root / "omni_features/local_qwen_1s" for key in ("data_root", "manifest", "feature_root", "target_root", "local_feature_root", "checkpoint_root", "result_root"): setattr(args, key, getattr(args, key).expanduser().resolve()) if len(set(args.variants)) != len(args.variants) or len(set(args.seeds)) != len(args.seeds): p.error("Duplicate variants or seeds") if any(seed < 0 for seed in args.seeds) or args.local_projection_dim < 0: p.error("Seeds and local-projection-dim must be nonnegative") if args.feature_root == args.local_feature_root: p.error("Global and local feature roots must be different") return args def run_paths(args, variant, seed): name = f"{variant}_seed{seed}" return args.checkpoint_root / f"{name}.pth", args.result_root / name def preflight(args) -> str: train = load_ego4d_manifest_split(args.manifest, "train") test = load_ego4d_manifest_split(args.manifest, "test") names = train + test print(f"EgoSpeak manifest: {len(train)} train / {len(test)} test; dev=none; epochs=4", flush=True) 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] if use_local: 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: for seed in args.seeds: checkpoint, result = run_paths(args, variant, seed) if checkpoint.exists() or checkpoint.with_suffix(".pth.json").exists() or result.exists(): raise FileExistsError(f"Refusing to overwrite {checkpoint} or {result}; " "select unfinished variants/seeds or new output roots") shape, contract = None, None counts = {"visual": [0, 0], "audio": [0, 0]} total = 0 noncanonical_target_rows = 0 target_row_pattern_counts: dict[str, int] = {} for index, name in enumerate(names, 1): data = load_session_data( args.feature_root, args.target_root, name, dataset="ego4d", expected_mode=args.expected_mode, local_feature_root=args.local_feature_root if use_local else None, local_modalities=local_modalities, local_contract=contract, ) if len(data.states) <= 10: raise ValueError(f"{name}: no training rows with all ten future horizons") current_shape = tuple(data.states.shape[1:]) if shape is not None and current_shape != shape: raise ValueError(f"Global feature shape differs: {name}: {current_shape} vs {shape}") shape = current_shape contract = data.feature_metadata.get("local_contract") if contract is not None: offset = 0 for modality in ("visual", "audio"): if local_modalities not in (modality, "av"): continue width = contract["config"][f"{modality}_dim"] counts[modality][0] += int(data.local_features[:, offset + width].sum()) counts[modality][1] += len(data.states) offset += width + 1 total += len(data.states) noncanonical_target_rows += int( data.feature_metadata.get("num_noncanonical_target_rows", 0) ) for pattern, count in data.feature_metadata.get( "target_row_pattern_counts", {} ).items(): target_row_pattern_counts[pattern] = ( target_row_pattern_counts.get(pattern, 0) + int(count) ) del data # Validate one clip at a time, including all test clips. if not args.no_progress and (index == 1 or index % 25 == 0 or index == len(names)): print(f"preflight: {index}/{len(names)} clips", flush=True) fusion = args.layer_fusion if fusion == "auto": fusion = "concat-mlp" if len(shape) == 2 else "single" if fusion == "single" and len(shape) != 1: raise ValueError("--layer-fusion single requires [T,D], not stacked global states") if fusion != "single" and len(shape) != 2: raise ValueError(f"--layer-fusion {fusion} requires stacked [T,L,D] global states") print(f"Validated {total} states; global shape per row={shape}; fusion={fusion}", flush=True) if noncanonical_target_rows: canonical_patterns = {"100", "010", "001", "011"} details = ", ".join( f"[{','.join(pattern)}]={count}" for pattern, count in sorted(target_row_pattern_counts.items()) if pattern not in canonical_patterns ) print( "Ego4D binary target rows outside the EasyCom four-state " f"constraint: {noncanonical_target_rows}/{total} ({details}); " "preserved exactly for the official LSTR MCE objective.", flush=True, ) for modality, (valid, count) in counts.items(): if count: note = "; provided boxes, not an online detector" if modality == "visual" else "" print(f"Local {modality} coverage: {valid / count:.2%} ({valid}/{count}){note}", flush=True) print("Training loads the full train split into host RAM; stacked global features may require tens of GB.", flush=True) return fusion def experiment_commands(args, fusion): common = ["--dataset", "ego4d", "--manifest", str(args.manifest), "--feature-root", str(args.feature_root), "--target-root", str(args.target_root), "--expected-mode", args.expected_mode, "--device", args.device] if args.no_progress: common.append("--no-progress") for variant in args.variants: for seed in args.seeds: checkpoint, result = run_paths(args, variant, seed) train = [sys.executable, str(REPO / "train_turn_head.py"), *common, "--output", str(checkpoint), "--epochs", "4", "--batch-size", "512", "--learning-rate", "7e-5", "--weight-decay", "5e-5", "--warmup-epochs", "1.6", "--warmup-factor", "0.3", "--dropout", "0.2", "--loss", "lstr-mce", "--transition-objective", "none", "--seed", str(seed)] if fusion != "single": # Unspecified means single in train_turn_head.py. train += ["--layer-fusion", fusion] infer = [sys.executable, str(REPO / "infer_turn_head.py"), *common, "--checkpoint", str(checkpoint), "--output-dir", str(result)] if variant != "baseline": train += ["--local-feature-root", str(args.local_feature_root), "--local-modalities", variant, "--local-projection-dim", str(args.local_projection_dim)] infer += ["--local-feature-root", str(args.local_feature_root)] yield f"{variant}_seed{seed}", train, infer def main(argv=None): args = parse_args(argv) fusion = preflight(args) env = os.environ.copy() for name in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"): env.setdefault(name, "4") for name, train, infer in experiment_commands(args, fusion): print(f"\n{name}: {'dry run' if args.dry_run else 'train, then evaluate'}", flush=True) for command in (train, infer): print(shlex.join(command), flush=True) if not args.dry_run: subprocess.run(command, cwd=REPO, env=env, check=True) summary = [sys.executable, str(REPO / "summarize_local_baselines.py"), "--root", str(args.result_root), "--dataset", "ego4d", "--paired-deltas"] print("\n" + shlex.join(summary), flush=True) if not args.dry_run: subprocess.run(summary, cwd=REPO, env=env, check=True) if __name__ == "__main__": main()