File size: 5,265 Bytes
72a99b3 3bea18f 72a99b3 3bea18f 05558cb 3bea18f 05558cb 72a99b3 05558cb 72a99b3 05558cb 72a99b3 05558cb 72a99b3 5bc0453 72a99b3 | 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 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 | """Evaluate one checkpoint across the full sweep grid and dump JSON.
Sweeps (fresh random mappings everywhere; nothing from training is reused):
- accuracy vs hop count k in 1..64 at the model's training-time unroll count
- accuracy vs inference-time unroll count n_iters in 1..16, for
k in {1, 2, 4, 8, 16, 32, 64} (both architectures: a 'fixed' model run
at n_iters>1 composes its stack repeatedly, i.e. deeper than trained)
- accuracy vs state-space size S' in {16, 32, 64} across k in 1..64
- LookupSolver reference: must be 1.0 (pipeline integrity check)
- inference FLOPs per example recorded per setting
Usage: python evaluate.py --ckpt out/rec1x4-seed0.pt --out-dir results
"""
import argparse
import json
import os
import numpy as np
import torch
from ptr.common import LookupSolver, sample_eval_batch
from ptr.flops import step_flops
from ptr.models import PtrTransformer
KS_FULL = list(range(1, 65))
KS_ITERS = [1, 2, 4, 8, 16, 32, 64]
ITERS_GRID = list(range(1, 17))
STATE_SIZES = [16, 32, 64]
def load_model(ckpt_path, device):
ck = torch.load(ckpt_path, map_location=device, weights_only=False)
cfg = ck["config"]
model = PtrTransformer(cfg["S"], cfg["d_model"], 4, 4 * cfg["d_model"],
depth=cfg["depth"],
stem_depth=cfg.get("stem_depth", 0))
model.load_state_dict(ck["state_dict"])
model.to(device).eval()
return model, ck
@torch.no_grad()
def acc_at_k(model, rng, S, k, n, n_iters, device):
b = sample_eval_batch(rng, n, S, k=k)
f = torch.as_tensor(b["f"], device=device)
x = torch.as_tensor(b["x"], device=device)
kt = torch.as_tensor(b["k"], device=device)
y = torch.as_tensor(b["answer"], device=device)
logits = model(f, x, kt, n_iters=n_iters)
# the head is sized to the trained S; for a smaller eval state space
# mask the logits to the first S classes (answers are always < S)
logits = logits[:, :S]
return float((logits.argmax(-1) == y).float().mean())
def evaluate_ckpt(ckpt_path, out_dir, n_per_cell=1000, seed=1234):
device = "cuda" if torch.cuda.is_available() else "cpu"
rng = np.random.default_rng(seed)
model, ck = load_model(ckpt_path, device)
name = os.path.splitext(os.path.basename(ckpt_path))[0]
cfg = ck["config"]
spec = ck["n_iters_spec"]
aligned = isinstance(spec, tuple) and spec[0] in (
"aligned", "aligned_margin")
margin = spec[3] if isinstance(spec, tuple) and len(spec) > 3 else 0
if isinstance(spec, tuple):
train_iters = (spec[1] + spec[2]) / 2 + margin
else:
train_iters = spec
def it_for(k):
# aligned variants read out at their trained offset: R = k + margin
return k + margin if aligned else int(train_iters)
is_rec = cfg["arch"] == "rec"
res = {"run_name": name, "config": cfg, "budget": ck["budget"],
"device": device, "n_per_cell": n_per_cell, "aligned": aligned}
# pipeline integrity: the lookup reference must be exact
b = sample_eval_batch(rng, 2000, cfg["S"], k=64)
assert (LookupSolver().solve(b["f"], b["x"], 64) == b["answer"]).all()
res["lookup_reference_acc"] = 1.0
# accuracy vs k at training-time unrolls (R=k for aligned models)
acc_k = {k: acc_at_k(model, rng, cfg["S"], k, n_per_cell, it_for(k),
device) for k in KS_FULL}
res["acc_vs_k"] = acc_k
res["train_iters"] = train_iters
res["inference_flops_per_example"] = step_flops(
model, cfg["S"], 1, train_iters, training=False)
# accuracy vs inference-time unroll count (OOD iterations)
iters_grid = ITERS_GRID
acc_iters = {str(r): {k: acc_at_k(model, rng, cfg["S"], k, n_per_cell, r,
device) for k in KS_ITERS}
for r in iters_grid}
res["acc_vs_iters"] = acc_iters
# accuracy vs state-space size (smaller spaces reuse trained vocab rows)
acc_state = {}
for Sp in STATE_SIZES:
if Sp > cfg["S"]:
continue
acc_state[str(Sp)] = {k: acc_at_k(model, rng, Sp, k, n_per_cell,
it_for(k), device) for k in
KS_FULL}
res["acc_vs_state_size"] = acc_state
os.makedirs(out_dir, exist_ok=True)
out_path = os.path.join(out_dir, f"{name}_eval.json")
with open(out_path, "w") as fh:
json.dump(res, fh)
hub = os.environ.get("HUB_REPO")
if hub:
from huggingface_hub import HfApi
HfApi().upload_file(path_or_fileobj=out_path,
path_in_repo=f"results/{name}_eval.json",
repo_id=hub, repo_type="model")
print(f"[hub] uploaded eval json to {hub}")
id64 = sum(v for k, v in acc_k.items() if k <= 8) / 8
print(f"[eval] {name}: mean acc k=1..8 {id64:.3f}, "
f"k=64 {acc_k[64]:.3f} (train_iters={train_iters}) -> {out_path}")
return out_path
if __name__ == "__main__":
p = argparse.ArgumentParser()
p.add_argument("--ckpt", required=True)
p.add_argument("--out-dir", default="results")
p.add_argument("--n-per-cell", type=int, default=1000)
a = p.parse_args()
evaluate_ckpt(a.ckpt, a.out_dir, a.n_per_cell) |