File size: 3,152 Bytes
ae02605 | 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 | """Does deep supervision make the recurrence EXTRAPOLATE past its training depth?
The claim behind supervising every recurrence is that it forces each application
of the shared block to be a valid one-step refinement. If that holds, the map is
iterable beyond where it was fit, and running more recurrences at inference than
at training should keep helping - or at least not collapse.
This is free: no retraining, only re-evaluating trained checkpoints at different
recurrence counts. A model whose accuracy peaks exactly at its training depth and
degrades past it has learned one entangled R-step function rather than a
repeatable refinement step.
"""
import argparse, glob, json, os, sys
from collections import defaultdict
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="rec1-*")
p.add_argument("--passes", type=int, default=16)
p.add_argument("--eval_bs", type=int, default=250)
p.add_argument("--out", default="figures/sudoku_rsweep.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
res = defaultdict(dict)
for d in sorted(glob.glob(os.path.join(a.runs, a.pattern))):
rf = os.path.join(d, "result.json")
if not (os.path.exists(rf) and os.path.exists(os.path.join(d, "model.pt"))):
continue
cfg = json.load(open(rf))["args"]
if not cfg.get("recurrent"):
continue
m = RecurrentDenoiser(len(tok), cfg["d"], cfg["layers"], cfg["heads"], cfg["pe"],
max_len=CELLS + 8, recurrences=cfg["R"]).to(dev)
m.load_state_dict(torch.load(os.path.join(d, "model.pt"), map_location=dev))
m.eval()
name = os.path.basename(d)
trained_R = cfg["R"]
print(f"\n--- {name} (trained with R={trained_R})", flush=True)
print(f" {'R at inference':>16}{'board':>10}", flush=True)
for R in (1, 2, 4, 8, 16, 32, 64, 128):
m.recurrences = R
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, a.passes, fillable=bl)
preds.append(y.cpu().numpy() + 1)
acc = float((np.concatenate(preds) == S).all(axis=1).mean())
res[name][R] = acc
flag = " <- trained here" if R == trained_R else ""
print(f" {R:>16}{acc*100:>9.2f}%{flag}", flush=True)
del m
torch.cuda.empty_cache()
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}")
|