| """Probe every final export under a run root, in parallel across GPUs. |
| |
| Reports AS-20K probe mAP against the patching's own random-init floor. That |
| ratio -- not reconstruction loss -- is the gate. This project has twice been |
| burned by treating MAE val_loss as a quality signal: cluster-weight-channels.md |
| measures it at rho=0.72 with the train/eval distribution-match channel, and |
| regmix-regression-256.md found it uncorrelated with downstream score at |
| rho=-0.034 over 256 mixtures. |
| """ |
| import argparse |
| import json |
| import queue |
| import subprocess |
| import threading |
| from pathlib import Path |
|
|
| REPO = "/workspace/code/eat-map-regmix" |
| FLOORS = {"p16": 0.03861, "p4x64": 0.04867} |
|
|
|
|
| def latest_export(run: Path) -> Path | None: |
| exports = sorted((run / "exports").glob("step_*")) if (run / "exports").exists() else [] |
| return exports[-1] if exports else None |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("root") |
| ap.add_argument("--gpus", default="0,1,2,3") |
| ap.add_argument("--all-steps", action="store_true", help="probe every export, not just the last") |
| |
| |
| |
| |
| |
| ap.add_argument("--per-gpu", type=int, default=2, help="concurrent probes per GPU") |
| ap.add_argument("--threads", type=int, default=12, help="BLAS threads per probe") |
| args = ap.parse_args() |
|
|
| jobs: queue.Queue = queue.Queue() |
| for run in sorted(Path(args.root).iterdir()): |
| if not run.is_dir(): |
| continue |
| exports = (sorted((run / "exports").glob("step_*")) if args.all_steps |
| else [latest_export(run)]) |
| for exp in [e for e in exports if e]: |
| if not (exp / "probe.json").exists(): |
| jobs.put((run.name, exp)) |
|
|
| total = jobs.qsize() |
| print(f"{total} exports to probe") |
| lock = threading.Lock() |
|
|
| def worker(gpu): |
| while True: |
| try: |
| name, exp = jobs.get_nowait() |
| except queue.Empty: |
| return |
| threads = str(args.threads) |
| proc = subprocess.run( |
| ["python", "-m", "eatmap.cli.probe", "--export", str(exp)], |
| cwd=REPO, capture_output=True, |
| env={**__import__("os").environ, "CUDA_VISIBLE_DEVICES": str(gpu), |
| "PYTHONPATH": REPO, |
| "OMP_NUM_THREADS": threads, "OPENBLAS_NUM_THREADS": threads, |
| "MKL_NUM_THREADS": threads, "NUMEXPR_NUM_THREADS": threads}, |
| ) |
| with lock: |
| if proc.returncode: |
| print(f" FAIL {name} {exp.name}: {proc.stderr.decode()[-200:]}") |
| jobs.task_done() |
|
|
| threads = [threading.Thread(target=worker, args=(int(g),)) |
| for g in args.gpus.split(",") for _ in range(args.per_gpu)] |
| for t in threads: |
| t.start() |
| for t in threads: |
| t.join() |
|
|
| rows = [] |
| for run in sorted(Path(args.root).iterdir()): |
| if not run.is_dir(): |
| continue |
| for exp in sorted((run / "exports").glob("step_*")): |
| pj = exp / "probe.json" |
| if not pj.exists(): |
| continue |
| d = json.loads(pj.read_text()) |
| tag = "p4x64" if "p4x64" in run.name else "p16" |
| floor = FLOORS[tag] |
| status = json.loads((run / "status.json").read_text()) if (run / "status.json").exists() else {} |
| rows.append((run.name, int(exp.name.split("_")[1]), d["probe/map"], |
| d["probe/map"] / floor, status.get("val/loss", float("nan")))) |
|
|
| print(f"\n{'run':<24}{'step':>8}{'probe mAP':>12}{'x floor':>10}{'val loss':>10}") |
| for name, step, m, ratio, vl in sorted(rows, key=lambda r: -r[2]): |
| print(f"{name:<24}{step:>8}{m:>12.5f}{ratio:>10.2f}{vl:>10.4f}") |
| Path(args.root, "probe_summary.json").write_text(json.dumps( |
| [{"run": n, "step": s, "probe_map": m, "x_floor": r, "val_loss": v} |
| for n, s, m, r, v in rows], indent=1)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
