File size: 4,831 Bytes
37fb52a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 | """Evaluate timed-out runs from their checkpoints.
The R-MDM comparison runs hit their wall-clock limit at step 80,000 of 110,000,
so no result.json was written. All three stopped at the SAME step, which keeps
the comparison internally valid - it is simply at 80k steps rather than 110k, and
that is stated wherever the numbers are used.
Configuration is reconstructed from the run name rather than a config file,
because the config is only written at the end of a successful run.
"""
import argparse, glob, json, os, sys
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from model import RecurrentDenoiser
import diffusion as dfn
from sudoku import CELLS, SudokuTokenizer
p = argparse.ArgumentParser()
p.add_argument("--hard", default="data/sudoku_hard")
p.add_argument("--runs", default="runs")
p.add_argument("--pattern", default="rmdm-*")
p.add_argument("--eval_bs", type=int, default=200)
p.add_argument("--out", default="figures/rmdm_compare.json")
a = p.parse_args()
dev = "cuda" if torch.cuda.is_available() else "cpu"
tok = SudokuTokenizer()
P = np.load(os.path.join(a.hard, "test_puz.npy"))
S = np.load(os.path.join(a.hard, "test_sol.npy"))
Pt, St = torch.from_numpy(P), torch.from_numpy(S - 1)
blank = Pt == 0
print(f"test set: {len(P)} puzzles, clues {(P>0).sum(1).min()}-{(P>0).sum(1).max()} "
f"(mean {(P>0).sum(1).mean():.1f})\n", flush=True)
def cfg_from_name(n):
"""All three runs share d=128, 1 layer, 4 heads, R=32 train / 64 infer.
They differ only in whether input injection is on."""
return {"d": 128, "layers": 1, "heads": 4, "pe": "ape", "R": 32, "R_infer": 64,
"step_embed": True, "inject": "faithful" not in n}
res = {}
for d in sorted(glob.glob(os.path.join(a.runs, a.pattern))):
ck = os.path.join(d, "ckpt.pt")
if not os.path.exists(ck):
continue
name = os.path.basename(d)
c = cfg_from_name(name)
st = torch.load(ck, map_location=dev)
m = RecurrentDenoiser(len(tok), c["d"], c["layers"], c["heads"], c["pe"],
max_len=CELLS + 8, recurrences=c["R_infer"],
inject=c["inject"], step_embed=c["step_embed"]).to(dev)
m.load_state_dict(st["model"])
m.eval()
print(f"--- {name} step {st['step']:,} injection={c['inject']} "
f"step_embed={c['step_embed']} params {m.n_params()/1e3:.0f}k", flush=True)
row = {"step": int(st["step"]), "inject": c["inject"],
"step_embed": c["step_embed"], "decode": {}}
print(f" {'passes':>8}{'board':>10}{'cell':>9}", flush=True)
for K in (1, 4, 8, 16, 32, 61):
preds = []
with torch.no_grad():
for i in range(0, len(Pt), a.eval_bs):
gold = St[i:i + a.eval_bs].to(dev)
bl = blank[i:i + a.eval_bs].to(dev)
x = torch.where(bl, torch.full_like(gold, tok.mask), gold)
y, _ = dfn.fixed_k_decode(m, x, tok, dev, K, fillable=bl)
preds.append(y.cpu().numpy() + 1)
pred = np.concatenate(preds)
board = float((pred == S).all(axis=1).mean())
cell = float((pred[P == 0] == S[P == 0]).mean())
row["decode"][str(K)] = {"board": board, "cell": cell}
print(f" {K:>8}{board*100:>9.2f}%{cell*100:>8.2f}%", flush=True)
row["best"] = max(v["board"] for v in row["decode"].values())
res[name] = row
del m
torch.cuda.empty_cache()
print(flush=True)
if res:
print("=" * 72)
print("HEAD-TO-HEAD: does input injection help once R-MDM's step embedding is present?")
print("=" * 72)
off = [v["best"] for v in res.values() if not v["inject"]]
on = [v["best"] for v in res.values() if v["inject"]]
print(f" R-MDM faithful (injection OFF) : {np.mean(off)*100:6.2f}% "
f"n={len(off)} {[round(x*100,2) for x in off]}")
if on:
print(f" + input injection : {np.mean(on)*100:6.2f}% "
f"n={len(on)} {[round(x*100,2) for x in on]}")
d = (np.mean(on) - np.mean(off)) * 100
print(f"\n input injection is worth {d:+.2f} points here")
if d <= 0.5:
print(" >>> Injection does NOT help once the step embedding is present.")
print(" >>> The earlier +5.60 was confounded with batch size and with the")
print(" >>> absence of R-MDM's per-loop conditioning. This must be reported.")
else:
print(" >>> Injection still helps with their conditioning present.")
print("\n All runs stopped at step 80,000 of 110,000 (wall-clock limit), so they")
print(" are matched to each other but undertrained relative to the headline runs.")
os.makedirs(os.path.dirname(a.out) or ".", exist_ok=True)
json.dump(res, open(a.out, "w"), indent=2)
print(f"\nwrote {a.out}")
|