Add Ego4D local AV ablation workflow
This commit is contained in:
Xin Wang
@@ -0,0 +1,175 @@
|
||||
"""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
|
||||
|
||||
|
||||
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="Your EXISTING global Omni features")
|
||||
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, ".pt"), (args.target_root, ".npy")]
|
||||
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()]
|
||||
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
|
||||
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)
|
||||
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)
|
||||
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()
|
||||
Reference in New Issue
Block a user