generated from mathias/template-go-web
feat(hpo): env-var knob overrides + sweep script (18 configs)
- 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>
This commit is contained in:
@@ -0,0 +1,97 @@
|
||||
"""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()
|
||||
Reference in New Issue
Block a user