generated from mathias/template-go-web
mathiasandClaude Sonnet 4.6
e739f84afd
- train.py knobs all readable from JEPA_* env vars (JEPA_WINDOW, JEPA_D_MODEL, JEPA_DEPTH, etc.) so hpo_sweep.py can override without touching source - scripts/hpo_sweep.py: 3×2×3 grid over D_MODEL × DEPTH × WINDOW, logs to results/hpo/hpo_results.jsonl with leaderboard at end - 3 new tests: env override correctness, configs() schema validation - 19/19 tests pass Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
98 lines
3.2 KiB
Python
98 lines
3.2 KiB
Python
"""HPO sweep for jepa-fx-risk HEPA backbone.
|
|
|
|
Runs train.py with different JEPA_* env overrides, logs results to
|
|
results/hpo/hpo_results.jsonl. Each config writes its metrics.json then
|
|
the result is appended to the JSONL.
|
|
|
|
Usage:
|
|
python scripts/hpo_sweep.py
|
|
python scripts/hpo_sweep.py --dry-run # print configs, don't train
|
|
"""
|
|
import argparse
|
|
import json
|
|
import os
|
|
import subprocess
|
|
import sys
|
|
from datetime import datetime
|
|
from itertools import product
|
|
from pathlib import Path
|
|
|
|
# ── Search space ──────────────────────────────────────────────────────────────
|
|
SEARCH_SPACE = {
|
|
"JEPA_D_MODEL": [64, 128, 256],
|
|
"JEPA_DEPTH": [2, 4],
|
|
"JEPA_WINDOW": [120, 240, 480],
|
|
}
|
|
# Fixed: PATCH_LEN=24 (1-day patches), N_HEADS=4, EPOCHS=300, PHASE1_EPOCHS=200
|
|
|
|
PYTHON = str(Path(sys.executable))
|
|
OUT_DIR = Path("results/hpo")
|
|
|
|
|
|
def configs():
|
|
"""Yield all configs as dicts of JEPA_* env overrides."""
|
|
keys = list(SEARCH_SPACE.keys())
|
|
for vals in product(*SEARCH_SPACE.values()):
|
|
yield dict(zip(keys, vals))
|
|
|
|
|
|
def run_config(cfg: dict, metrics_path: str = "metrics.json") -> dict:
|
|
env = {**os.environ, **{k: str(v) for k, v in cfg.items()}}
|
|
result = subprocess.run(
|
|
[PYTHON, "train.py"],
|
|
env=env,
|
|
capture_output=True,
|
|
text=True,
|
|
)
|
|
if result.returncode != 0:
|
|
return {"config": cfg, "error": result.stderr[-500:]}
|
|
stdout_last = result.stdout.strip().split("\n")[-1]
|
|
with open(metrics_path) as f:
|
|
m = json.load(f)
|
|
return {
|
|
"config": cfg,
|
|
"val_vol_r2": m.get("val_vol_r2"),
|
|
"phase1_r2": m.get("phase1_r2"),
|
|
"stdout_last": stdout_last,
|
|
}
|
|
|
|
|
|
def main():
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument("--dry-run", action="store_true")
|
|
args = parser.parse_args()
|
|
|
|
OUT_DIR.mkdir(parents=True, exist_ok=True)
|
|
out_file = OUT_DIR / "hpo_results.jsonl"
|
|
|
|
all_cfgs = list(configs())
|
|
print(f"HPO sweep: {len(all_cfgs)} configs")
|
|
for i, cfg in enumerate(all_cfgs):
|
|
label = " ".join(f"{k.replace('JEPA_','')}={v}" for k, v in cfg.items())
|
|
print(f"\n[{i+1}/{len(all_cfgs)}] {label}")
|
|
if args.dry_run:
|
|
continue
|
|
ts = datetime.utcnow().isoformat()
|
|
row = run_config(cfg)
|
|
row["ts"] = ts
|
|
with open(out_file, "a") as f:
|
|
f.write(json.dumps(row) + "\n")
|
|
if "error" in row:
|
|
print(f" ERROR: {row['error'][:200]}")
|
|
else:
|
|
print(f" val_vol_r2={row['val_vol_r2']:.4f} phase1_r2={row['phase1_r2']:.4f}")
|
|
|
|
if not args.dry_run:
|
|
# Print leaderboard
|
|
rows = [json.loads(l) for l in open(out_file) if l.strip()]
|
|
rows = [r for r in rows if "error" not in r]
|
|
rows.sort(key=lambda r: r.get("phase1_r2", -999), reverse=True)
|
|
print("\n── Leaderboard (by phase1_r2) ─────────────────────────")
|
|
for r in rows[:5]:
|
|
cfg_str = " ".join(f"{k.replace('JEPA_','')}={v}" for k,v in r["config"].items())
|
|
print(f" {r['phase1_r2']:.4f} {cfg_str}")
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|