134 lines
6.8 KiB
Python
134 lines
6.8 KiB
Python
|
"""Run raw encoder-only Ego4D heads: four epochs, no dev, manifest splits.
|
||
|
|
|
||
|
|
This is not a capacity/context-matched Thinker ablation: native encoder vectors
|
||
|
|
and validity flags are concatenated, so the existing MLP adapts to input width.
|
||
|
|
"""
|
||
|
|
|
||
|
|
from __future__ import annotations
|
||
|
|
|
||
|
|
import argparse
|
||
|
|
import os
|
||
|
|
from pathlib import Path
|
||
|
|
import shlex
|
||
|
|
import subprocess
|
||
|
|
import sys
|
||
|
|
|
||
|
|
from tqdm.auto import tqdm
|
||
|
|
|
||
|
|
from easycom_head import load_ego4d_manifest_split, load_session_data
|
||
|
|
from local_turn_features import LOCAL_FORMAT, MODALITIES
|
||
|
|
from prepare_encoder_only_features import DEFAULT_DATA_ROOT, PACKING, encoder_only_mode
|
||
|
|
|
||
|
|
|
||
|
|
REPO = Path(__file__).resolve().parent
|
||
|
|
|
||
|
|
|
||
|
|
def parse_args(argv=None):
|
||
|
|
parser = argparse.ArgumentParser(description=__doc__)
|
||
|
|
parser.add_argument("--data-root", type=Path, default=DEFAULT_DATA_ROOT)
|
||
|
|
parser.add_argument("--manifest", type=Path)
|
||
|
|
parser.add_argument("--target-root", type=Path)
|
||
|
|
parser.add_argument("--feature-root", type=Path, required=True, help="Prepared encoder-only inputs, not Thinker/local sidecars")
|
||
|
|
parser.add_argument("--modalities", choices=tuple(MODALITIES), default="av")
|
||
|
|
parser.add_argument("--checkpoint-root", type=Path, default=Path("checkpoints/ego4d_encoder_only_1s"))
|
||
|
|
parser.add_argument("--result-root", type=Path, default=Path("results/ego4d_encoder_only_1s"))
|
||
|
|
parser.add_argument("--seeds", type=int, nargs="+", default=[0, 1, 2])
|
||
|
|
parser.add_argument("--device", default="cuda:0")
|
||
|
|
parser.add_argument("--dry-run", action="store_true", help="Validate all train/test inputs and print commands without training")
|
||
|
|
parser.add_argument("--no-progress", action="store_true")
|
||
|
|
args = parser.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"
|
||
|
|
for key in ("data_root", "manifest", "target_root", "feature_root", "checkpoint_root", "result_root"):
|
||
|
|
setattr(args, key, getattr(args, key).expanduser().resolve())
|
||
|
|
if len(set(args.seeds)) != len(args.seeds) or any(seed < 0 for seed in args.seeds):
|
||
|
|
parser.error("Seeds must be unique and nonnegative")
|
||
|
|
return args
|
||
|
|
|
||
|
|
|
||
|
|
def run_paths(args, seed):
|
||
|
|
name = f"encoder_{args.modalities}_seed{seed}"
|
||
|
|
return args.checkpoint_root / f"{name}.pth", args.result_root / name
|
||
|
|
|
||
|
|
|
||
|
|
def preflight(args):
|
||
|
|
names = load_ego4d_manifest_split(args.manifest, "all")
|
||
|
|
train = load_ego4d_manifest_split(args.manifest, "train")
|
||
|
|
test = load_ego4d_manifest_split(args.manifest, "test")
|
||
|
|
marker = args.feature_root / "complete.txt"
|
||
|
|
if not marker.is_file() or marker.read_text().strip() != f"{len(names)} clips":
|
||
|
|
raise ValueError(f"Missing/incomplete encoder-only export: {marker}")
|
||
|
|
for seed in args.seeds:
|
||
|
|
checkpoint, result = run_paths(args, seed)
|
||
|
|
if checkpoint.exists() or checkpoint.with_suffix(".pth.json").exists() or result.exists():
|
||
|
|
raise FileExistsError(f"Refusing to overwrite {checkpoint} or {result}; choose unfinished seeds/new output roots")
|
||
|
|
expected, total, width = None, 0, None
|
||
|
|
for name in tqdm(names, desc="encoder-only inputs", unit="clip", disable=args.no_progress):
|
||
|
|
data = load_session_data(args.feature_root, args.target_root, name, dataset="ego4d",
|
||
|
|
expected_mode=encoder_only_mode(args.modalities))
|
||
|
|
config = data.feature_metadata.get("config") or {}
|
||
|
|
contract = config.get("encoder_contract") or {}
|
||
|
|
if (config.get("input_source") != "qwen-encoders-only" or config.get("uses_thinker_states") is not False
|
||
|
|
or config.get("packing") != PACKING or contract.get("format") != LOCAL_FORMAT
|
||
|
|
or contract.get("modalities") != args.modalities):
|
||
|
|
raise ValueError(f"{name}: not a matching encoder-only export")
|
||
|
|
encoder_config = contract.get("config") or {}
|
||
|
|
if (encoder_config.get("causal_policy") != "half-open-trailing-window"
|
||
|
|
or encoder_config.get("uses_speech_labels") is not False):
|
||
|
|
raise ValueError(f"{name}: missing causal/no-speech-label encoder contract")
|
||
|
|
dimensions = [encoder_config.get(f"{part}_dim") for part in MODALITIES[args.modalities]]
|
||
|
|
if any(not isinstance(d, int) or d <= 0 for d in dimensions):
|
||
|
|
raise ValueError(f"{name}: invalid encoder dimensions")
|
||
|
|
width = sum(d + 1 for d in dimensions)
|
||
|
|
if data.states.ndim != 2 or data.hidden_size != width or len(data.states) <= 10:
|
||
|
|
raise ValueError(f"{name}: invalid native encoder-vector/validity-bit shape")
|
||
|
|
if expected is not None and config != expected:
|
||
|
|
raise ValueError(f"{name}: encoder export contract differs across clips")
|
||
|
|
expected = config
|
||
|
|
total += len(data.states)
|
||
|
|
del data
|
||
|
|
print(f"Encoder-only {args.modalities}: {len(train)} train / {len(test)} test; "
|
||
|
|
f"{total} states, input width={width}, MLP hidden width={width // 2}; epochs=4, dev=none", flush=True)
|
||
|
|
print("No Thinker features or local-fusion branch. Native concatenation changes head capacity; "
|
||
|
|
"context remains the saved encoder windows.", flush=True)
|
||
|
|
|
||
|
|
|
||
|
|
def experiment_commands(args):
|
||
|
|
common = ["--dataset", "ego4d", "--manifest", str(args.manifest),
|
||
|
|
"--feature-root", str(args.feature_root), "--target-root", str(args.target_root),
|
||
|
|
"--expected-mode", encoder_only_mode(args.modalities), "--device", args.device]
|
||
|
|
if args.no_progress:
|
||
|
|
common.append("--no-progress")
|
||
|
|
for seed in args.seeds:
|
||
|
|
checkpoint, result = run_paths(args, 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)]
|
||
|
|
infer = [sys.executable, str(REPO / "infer_turn_head.py"), *common,
|
||
|
|
"--checkpoint", str(checkpoint), "--output-dir", str(result)]
|
||
|
|
yield train, infer
|
||
|
|
|
||
|
|
|
||
|
|
def main(argv=None):
|
||
|
|
args = parse_args(argv)
|
||
|
|
preflight(args)
|
||
|
|
env = os.environ.copy()
|
||
|
|
for name in ("OMP_NUM_THREADS", "MKL_NUM_THREADS", "OPENBLAS_NUM_THREADS"):
|
||
|
|
env.setdefault(name, "4")
|
||
|
|
for commands in experiment_commands(args):
|
||
|
|
for command in commands:
|
||
|
|
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"]
|
||
|
|
print(shlex.join(summary), flush=True)
|
||
|
|
if not args.dry_run:
|
||
|
|
subprocess.run(summary, cwd=REPO, env=env, check=True)
|
||
|
|
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
main()
|