summaryrefslogtreecommitdiff
path: root/ep_run/auto_probe.py
diff options
context:
space:
mode:
authorYuren Hao <yurenh2@illinois.edu>2026-07-03 05:56:50 -0500
committerYuren Hao <yurenh2@illinois.edu>2026-07-03 05:56:50 -0500
commitb83947778e2c776f757a07d4719b7ce961d7ed55 (patch)
treeb9cc01d7adda691d9156d9d04f4fb2f644674e96 /ep_run/auto_probe.py
Initial commit: ept — backprop-free equilibrium transformer (EP)
Code (ep_run/), organized docs (docs/{method,campaign,hardware,outreach,paper}), analysis scripts (scripts/), ONBOARDING.md entry point. Large data/checkpoints git-ignored (share separately). Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_014FAPDWQ49M5Ye3NpTndTpn
Diffstat (limited to 'ep_run/auto_probe.py')
-rw-r--r--ep_run/auto_probe.py25
1 files changed, 25 insertions, 0 deletions
diff --git a/ep_run/auto_probe.py b/ep_run/auto_probe.py
new file mode 100644
index 0000000..90b969e
--- /dev/null
+++ b/ep_run/auto_probe.py
@@ -0,0 +1,25 @@
+"""Wait for the first converged clean-EP ckpt, run the fixed oracle-adjoint probe on it, report g_EP vs g_transpose."""
+import time, os, subprocess, shutil
+WD = "/home/yurenh2/ept/ep_run"; os.chdir(WD)
+CK, FROZEN = "runs/ep_clean.pt", "runs/ep_clean_probe.pt"
+got = False
+for _ in range(45): # up to ~67 min
+ time.sleep(90)
+ if os.path.exists(CK) and os.path.getsize(CK) > 1_000_000:
+ got = True; break
+if not got:
+ print("=== AUTO-PROBE: ep_clean.pt never appeared in ~67min ==="); raise SystemExit
+shutil.copy2(CK, FROZEN) # freeze (avoid write race)
+env = dict(os.environ, CUDA_VISIBLE_DEVICES="0", PYTORCH_CUDA_ALLOC_CONF="expandable_segments:True")
+try:
+ r = subprocess.run(["python3", "asym_probe.py", "--ckpt", FROZEN, "--B", "8"],
+ env=env, capture_output=True, text=True, timeout=1200)
+ txt = r.stdout + "\n" + r.stderr
+except Exception as e:
+ txt = f"probe run error: {e}"
+print("=== AUTO ORACLE PROBE on first clean-EP ckpt ===")
+KEEP = ("z* res", "GMRES", "resid", "cos(", "g_transpose", "g_EP", "g_BPTT", "interpret", "->", "exact", "AsymEP", "# ckpt", "step ")
+DROP = ("UserWarning", "cuBLAS", "warnings.warn", "Triggered", "return Variable", "FutureWarning")
+for line in txt.splitlines():
+ if any(k in line for k in KEEP) and not any(b in line for b in DROP):
+ print(line)