icassp2027-experiment/run_ego4d_local_baselines.py

"""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()