Add hyperparameter scan
This commit is contained in:
@@ -0,0 +1,192 @@
|
||||
#!/usr/bin/env python3
|
||||
"""Hyperparameter scan over dropout x n_blocks x hidden_dim.
|
||||
|
||||
Runs `giant train` sequentially (this machine has a single GPU) for every
|
||||
combination, plus one extra run at the default architecture with a higher
|
||||
learning rate. Runs are shuffled so the parameter space gets coarse coverage
|
||||
early rather than exhausting one corner of the grid first.
|
||||
|
||||
Usage:
|
||||
uv run python scripts/hparam_scan.py
|
||||
uv run python scripts/hparam_scan.py --dry-run
|
||||
uv run python scripts/hparam_scan.py --seed 1 --data /path/to/parquet
|
||||
"""
|
||||
|
||||
import argparse
|
||||
import csv
|
||||
import itertools
|
||||
import os
|
||||
import random
|
||||
import subprocess
|
||||
import time
|
||||
from pathlib import Path
|
||||
|
||||
DATA_DEFAULT = "/home/lars/geant_steps/train"
|
||||
SCAN_DIR_DEFAULT = "checkpoints/scan"
|
||||
EPOCHS = 50
|
||||
DEFAULT_LR = 3e-4
|
||||
|
||||
DROPOUTS = [0.0, 0.1]
|
||||
N_BLOCKS = [5, 6, 7, 8]
|
||||
HIDDEN_DIMS = [256, 512, 1024]
|
||||
|
||||
EXTRA_LR_RUN = {"hidden_dim": 256, "n_blocks": 6, "dropout": 0.1, "lr": 1e-3}
|
||||
|
||||
_SUMMARY_FIELDS = [
|
||||
"name",
|
||||
"hidden_dim",
|
||||
"n_blocks",
|
||||
"dropout",
|
||||
"lr",
|
||||
"epochs_completed",
|
||||
"final_val_loss",
|
||||
"best_val_loss",
|
||||
"wall_time_s",
|
||||
]
|
||||
|
||||
|
||||
def build_runs(seed: int) -> list[dict]:
|
||||
runs = [
|
||||
{"hidden_dim": h, "n_blocks": n, "dropout": d, "lr": DEFAULT_LR}
|
||||
for d, n, h in itertools.product(DROPOUTS, N_BLOCKS, HIDDEN_DIMS)
|
||||
]
|
||||
runs.append(dict(EXTRA_LR_RUN))
|
||||
random.Random(seed).shuffle(runs)
|
||||
return runs
|
||||
|
||||
|
||||
def run_name(run: dict) -> str:
|
||||
return f"h{run['hidden_dim']}_n{run['n_blocks']}_d{run['dropout']}_lr{run['lr']}"
|
||||
|
||||
|
||||
def last_completed_epoch(metrics_path: Path) -> int:
|
||||
if not metrics_path.exists():
|
||||
return 0
|
||||
with open(metrics_path, newline="") as f:
|
||||
rows = list(csv.DictReader(f))
|
||||
if not rows:
|
||||
return 0
|
||||
return int(rows[-1]["epoch"])
|
||||
|
||||
|
||||
def final_metrics(metrics_path: Path) -> tuple[int, float, float]:
|
||||
with open(metrics_path, newline="") as f:
|
||||
rows = list(csv.DictReader(f))
|
||||
epochs_completed = int(rows[-1]["epoch"])
|
||||
final_val_loss = float(rows[-1]["val_loss"])
|
||||
best_val_loss = min(float(r["val_loss"]) for r in rows)
|
||||
return epochs_completed, final_val_loss, best_val_loss
|
||||
|
||||
|
||||
def append_summary(summary_path: Path, row: dict) -> None:
|
||||
write_header = not summary_path.exists()
|
||||
with open(summary_path, "a", newline="") as f:
|
||||
writer = csv.DictWriter(f, fieldnames=_SUMMARY_FIELDS)
|
||||
if write_header:
|
||||
writer.writeheader()
|
||||
writer.writerow(row)
|
||||
|
||||
|
||||
def main() -> None:
|
||||
parser = argparse.ArgumentParser(description=__doc__)
|
||||
parser.add_argument("--data", default=DATA_DEFAULT)
|
||||
parser.add_argument("--scan-dir", default=SCAN_DIR_DEFAULT)
|
||||
parser.add_argument("--seed", type=int, default=0)
|
||||
parser.add_argument("--dry-run", action="store_true")
|
||||
args = parser.parse_args()
|
||||
|
||||
runs = build_runs(args.seed)
|
||||
scan_dir = Path(args.scan_dir)
|
||||
|
||||
if args.dry_run:
|
||||
for i, run in enumerate(runs, 1):
|
||||
print(f"[{i}/{len(runs)}] {run_name(run)}")
|
||||
return
|
||||
|
||||
scan_dir.mkdir(parents=True, exist_ok=True)
|
||||
summary_path = scan_dir / "scan_summary.csv"
|
||||
|
||||
env = os.environ.copy()
|
||||
env["TQDM_DISABLE"] = "1"
|
||||
|
||||
for i, run in enumerate(runs, 1):
|
||||
name = run_name(run)
|
||||
out_dir = scan_dir / name
|
||||
metrics_path = out_dir / "metrics.csv"
|
||||
last_ckpt = out_dir / "last.pt"
|
||||
|
||||
completed = last_completed_epoch(metrics_path)
|
||||
if completed >= EPOCHS:
|
||||
print(f"[{i}/{len(runs)}] {name} — already complete, skipping")
|
||||
continue
|
||||
|
||||
out_dir.mkdir(parents=True, exist_ok=True)
|
||||
cmd = [
|
||||
"giant",
|
||||
"train",
|
||||
args.data,
|
||||
"--mode",
|
||||
"flow",
|
||||
"--epochs",
|
||||
str(EPOCHS),
|
||||
"--batch-size",
|
||||
"auto",
|
||||
"--hidden-dim",
|
||||
str(run["hidden_dim"]),
|
||||
"--n-blocks",
|
||||
str(run["n_blocks"]),
|
||||
"--dropout",
|
||||
str(run["dropout"]),
|
||||
"--lr",
|
||||
str(run["lr"]),
|
||||
"--out",
|
||||
str(out_dir),
|
||||
]
|
||||
if last_ckpt.exists():
|
||||
cmd += ["--resume", str(last_ckpt)]
|
||||
print(f"[{i}/{len(runs)}] {name} — resuming from epoch {completed}")
|
||||
else:
|
||||
print(f"[{i}/{len(runs)}] {name} — starting")
|
||||
|
||||
start = time.monotonic()
|
||||
try:
|
||||
with open(out_dir / "train.log", "a") as log:
|
||||
subprocess.run(cmd, env=env, stdout=log, stderr=subprocess.STDOUT)
|
||||
except KeyboardInterrupt:
|
||||
print(
|
||||
f"\ninterrupted during {name} — re-run this script to resume "
|
||||
f"(checkpoint/resume is handled by `giant train` itself)"
|
||||
)
|
||||
return
|
||||
wall_time_s = time.monotonic() - start
|
||||
|
||||
if metrics_path.exists():
|
||||
epochs_completed, final_val_loss, best_val_loss = final_metrics(
|
||||
metrics_path
|
||||
)
|
||||
append_summary(
|
||||
summary_path,
|
||||
{
|
||||
"name": name,
|
||||
"hidden_dim": run["hidden_dim"],
|
||||
"n_blocks": run["n_blocks"],
|
||||
"dropout": run["dropout"],
|
||||
"lr": run["lr"],
|
||||
"epochs_completed": epochs_completed,
|
||||
"final_val_loss": final_val_loss,
|
||||
"best_val_loss": best_val_loss,
|
||||
"wall_time_s": round(wall_time_s, 1),
|
||||
},
|
||||
)
|
||||
print(
|
||||
f"[{i}/{len(runs)}] {name} — val_loss {final_val_loss:.4f} "
|
||||
f"({wall_time_s:.1f}s)"
|
||||
)
|
||||
else:
|
||||
print(
|
||||
f"[{i}/{len(runs)}] {name} — no metrics.csv produced, check train.log"
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user