151 lines
7.7 KiB
Python
151 lines
7.7 KiB
Python
"""Run raw encoder-only heads: four epochs, no dev, fixed benchmark 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.
|
|
This entry defaults to Ego4D; run_easycom_encoder_only.py shares the workflow.
|
|
"""
|
|
|
|
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 DATASET_NAMES, load_session_data
|
|
from local_turn_features import LOCAL_FORMAT, MODALITIES
|
|
from prepare_encoder_only_features import (
|
|
DEFAULT_DATA_ROOTS, PACKING, encoder_only_mode, encoder_sample_split, sample_unit,
|
|
)
|
|
|
|
|
|
REPO = Path(__file__).resolve().parent
|
|
|
|
|
|
def parse_args(argv=None, *, dataset="ego4d"):
|
|
if dataset not in DATASET_NAMES:
|
|
raise ValueError(f"Unknown dataset: {dataset}")
|
|
parser = argparse.ArgumentParser(
|
|
description=f"{dataset}: raw encoder-only A/V/AV heads; four epochs, no dev, fixed train/test splits."
|
|
)
|
|
parser.add_argument("--data-root", type=Path, default=DEFAULT_DATA_ROOTS[dataset])
|
|
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(f"checkpoints/{dataset}_encoder_only_1s"))
|
|
parser.add_argument("--result-root", type=Path, default=Path(f"results/{dataset}_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.dataset = dataset
|
|
if dataset == "ego4d":
|
|
args.manifest = args.manifest or args.data_root / "ego4d_egospeak.tsv"
|
|
elif args.manifest is not None:
|
|
parser.error("EasyCom uses train sessions 4..12 and test sessions 3/2/1, not --manifest")
|
|
args.target_root = args.target_root or args.data_root / f"{dataset}_target_frame"
|
|
for key in ("data_root", "manifest", "target_root", "feature_root", "checkpoint_root", "result_root"):
|
|
value = getattr(args, key)
|
|
if value is not None:
|
|
setattr(args, key, value.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, train, test = encoder_sample_split(args.dataset, args.manifest)
|
|
unit = sample_unit(args.dataset)
|
|
marker = args.feature_root / "complete.txt"
|
|
if not marker.is_file() or marker.read_text().strip() != f"{len(names)} {unit}s":
|
|
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=unit, disable=args.no_progress):
|
|
data = load_session_data(args.feature_root, args.target_root, name, dataset=args.dataset,
|
|
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 {unit}s")
|
|
expected = config
|
|
total += len(data.states)
|
|
del data
|
|
print(f"{args.dataset} 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", args.dataset, "--feature-root", str(args.feature_root), "--target-root", str(args.target_root),
|
|
"--expected-mode", encoder_only_mode(args.modalities), "--device", args.device]
|
|
if args.manifest is not None:
|
|
common += ["--manifest", str(args.manifest)]
|
|
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)]
|
|
if args.dataset == "easycom":
|
|
_, train_sessions, test_sessions = encoder_sample_split("easycom")
|
|
train += ["--sessions", *train_sessions]
|
|
infer += ["--sessions", *test_sessions]
|
|
yield train, infer
|
|
|
|
|
|
def main(argv=None, *, dataset="ego4d"):
|
|
args = parse_args(argv, dataset=dataset)
|
|
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", args.dataset]
|
|
print(shlex.join(summary), flush=True)
|
|
if not args.dry_run:
|
|
subprocess.run(summary, cwd=REPO, env=env, check=True)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|