summaryrefslogtreecommitdiff
path: root/ep_run/newstop_gate.py
diff options
context:
space:
mode:
Diffstat (limited to 'ep_run/newstop_gate.py')
-rw-r--r--ep_run/newstop_gate.py32
1 files changed, 32 insertions, 0 deletions
diff --git a/ep_run/newstop_gate.py b/ep_run/newstop_gate.py
new file mode 100644
index 0000000..dc01e17
--- /dev/null
+++ b/ep_run/newstop_gate.py
@@ -0,0 +1,32 @@
+"""Ship-gate for holo_a_track_avg (trend-aware stop + plateau averaging, the semi-convergence fix):
+cos(EP,BPTT) at t2sel=160 on s2000 — the regime where argmin-t_best got fooled (t2_probe: batch2
+degraded 0.956->0.894 at 160). PASS = avg recovers the t2sel-80 level (~0.93+) at 160 without
+hurting the healthy batches."""
+import torch
+import lt_ep_train as L
+from diag_cos import cos_ep_bptt
+import holo_ep
+
+torch.manual_seed(0)
+blk = L.EQBlock(512, 16, 256, 256, c=1.0, attn_mode='thick'); blk.qknorm = True
+ck = torch.load('runs/redx_traj/s2000.pt', map_location=L.dev)
+with torch.no_grad():
+ for p, w in zip(blk.allp, ck['allp']):
+ p.copy_(w.to(L.dev))
+blk.track = True
+blk.holofast = True
+torch.manual_seed(11)
+batches = [L.get_batch('train', 24, 256) for _ in range(3)]
+
+orig_fast = holo_ep.holo_a_track_fast
+for name, fn in (('argmin(fast)', holo_ep.holo_a_track_fast), ('trend+avg', holo_ep.holo_a_track_avg)):
+ holo_ep.holo_a_track_fast = fn # ep_step routes track via holofast -> this symbol
+ import lt_ep_train
+ lt_ep_train.ep_step.__globals__ # (route happens at import inside ep_step)
+ cs = []
+ for idx, y in batches:
+ c, _ = cos_ep_bptt(blk, idx, y, 150, 20, 0.1, 0.02, holo=2, hr=0.02, t2sel=160, bsub=4)
+ cs.append(c)
+ print(f"{name:>13}: cos={' '.join(f'{c:.4f}' for c in cs)} mean={sum(cs)/len(cs):.4f}", flush=True)
+holo_ep.holo_a_track_fast = orig_fast
+print("NEWSTOP_GATE_DONE", flush=True)