| """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}") |
|
|
