File size: 6,743 Bytes
1cb3dc4 | 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 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 | """Score a saved checkpoint on the CE(q||p) <= log|S| promotion criterion.
Runs the same evaluation the trainer runs, on the same held-out split, but
offline against an arbitrary checkpoint directory, and reports the per-stage
value statistics with the CE criterion first.
The criterion, per cell with |S| >= 2 and headroom:
CE(q||p) <= log|S|
where q is the empirical digit distribution of that puzzle's stored instances at
that stage (post-constraint, what the data teaches) and log|S| is the entropy of
the uniform distribution over the raw pre-constraint candidate set. Cells where
q is already uniform on S are skipped: there the bar equals the floor H(q), so by
Gibbs only p == q exactly could pass.
Usage:
python wavecurriculum_run/eval_ceq_ckpt.py --ckpt_dir DIR [--stage 6]
[--eval_epochs 40] [--label NAME]
DIR is a directory containing a checkpoint_<step> file, as unpacked from the
latest.tar / stageN_step.tar archives on the Hub.
"""
import argparse
import functools
import json
import os
import sys
os.environ.setdefault("SUDOKU_VERBOSE_EVAL", "0")
import numpy as np
def main():
ap = argparse.ArgumentParser()
ap.add_argument("--ckpt_dir", required=True)
ap.add_argument("--stage", type=int, default=0,
help="Curriculum stage to evaluate at; 0 = read it from "
"the checkpoint's sibling latest.json, else max_stage.")
ap.add_argument("--eval_epochs", type=int, default=40)
ap.add_argument("--label", default="")
ap.add_argument("--out_json", default="")
args = ap.parse_args()
# The config is assembled from the environment exactly as the training jobs
# assemble it, so the launcher must export the same SUDOKU_* variables. Only
# the eval-side knobs are overridden here.
os.environ["SUDOKU_EVAL_EPOCHS"] = str(args.eval_epochs)
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import jax
import jax.numpy as jnp
from jax import random
from flax.training import checkpoints
from flax import jax_utils
from train import main as train_main
from train import data, model, trainer, evaluater
config = train_main.get_config()
K = int(config.num_latent_slots)
stage = args.stage or int(getattr(config, "curriculum_max_stage", 6))
# Eval iterator only; the curriculum stage matters because the candidate
# masks and instance pool are read per stage.
curriculum = data.CurriculumState(
stage=stage, max_stage=int(getattr(config, "curriculum_max_stage", 6)))
eval_iter = data.create_iter(config, config.minibatch_size, train=False,
curriculum=curriculum)
model_config = model.TransformerConfig(
dtype=config.dtype,
vocab_size=config.vocab_size,
seq_len=config.seq_len,
num_heads=config.num_heads,
num_layers=config.num_layers,
emb_dim=config.emb_dim,
qkv_dim=config.qkv_dim,
mlp_dim=config.mlp_dim,
dropout_rate=config.dropout_rate,
attention_dropout_rate=config.attention_dropout_rate,
deterministic=False,
num_latent_slots=K,
inject_latents=bool(int(getattr(config, "recurrent_latent", 1))),
)
rng = jax.random.PRNGKey(config.seed)
rng, init_rng, _ = random.split(rng, num=3)
rng, dropout_rng = jax.random.split(rng)
net = model.TransformerLMHeadModel(model_config)
dummy_latents = jnp.zeros(
(config.minibatch_size, K, config.emb_dim), model_config.dtype)
dummy_positions = jnp.zeros((config.minibatch_size, K), jnp.int32)
dummy_active = jnp.zeros((config.minibatch_size, K), bool)
_, initial_variables = jax.jit(net.init_with_output)(
{"params": init_rng, "dropout": dropout_rng},
jnp.ones((config.minibatch_size, config.seq_len), jnp.int32),
dummy_latents, dummy_positions, dummy_active)
state, _ = trainer.get_state(config, net, initial_variables)
restored = checkpoints.restore_checkpoint(os.path.abspath(args.ckpt_dir),
state)
step = int(restored.step)
if step == 0:
sys.exit(f"no checkpoint restored from {args.ckpt_dir}")
print(f"restored step {step} from {args.ckpt_dir}", flush=True)
state = jax_utils.replicate(restored)
p_eval_step = jax.pmap(
functools.partial(evaluater.eval_step,
config=model_config.replace(deterministic=True)),
axis_name="batch")
em = evaluater.get_eval_metrics(state, eval_iter, p_eval_step, config)
ceq = em.get("per_stage_val_ceq", {})
hq = em.get("per_stage_val_hq", {})
ok = em.get("per_stage_val_ceq_ok", {})
gap = em.get("per_stage_val_ceq_gap", {})
puz = em.get("per_stage_val_ceq_puz", {})
mass = em.get("per_stage_val_mass_multi", {})
spread = em.get("per_stage_val_spread", {})
hb = em.get("per_stage_val_hbound", {})
tag = args.label or os.path.basename(os.path.abspath(args.ckpt_dir))
print(f"\n==== {tag} step {step} eval_stage {stage} "
f"({args.eval_epochs} eval batches) ====")
print("stage H(q) ceq log|S|* slack ceq_ok ceq_puz"
" mass* spread* hbound")
rows = {}
for s in sorted(ceq):
if ceq[s] < 0:
continue
# log|S| on the counted cells, recovered from ceq + its slack.
lm = ceq[s] + gap[s]
rows[s] = {"h_q": hq[s], "ceq": ceq[s], "log_m": lm, "slack": gap[s],
"ceq_ok": ok[s], "ceq_puz": puz.get(s, -1.0),
"mass_multi": mass.get(s, -1.0),
"spread": spread.get(s, -1.0), "hbound": hb.get(s, -1.0)}
flag = "" if gap[s] >= 0 else " <-- over ceiling"
print(f"{s:>5} {hq[s]:6.4f} {ceq[s]:7.4f} {lm:7.4f} {gap[s]:+7.4f}"
f" {ok[s]:6.3f} {puz.get(s, -1):7.3f}"
f" {mass.get(s, -1):6.3f} {spread.get(s, -1):7.3f}"
f" {hb.get(s, -1):6.3f}{flag}")
if stage in rows:
r = rows[stage]
verdict = ("PROMOTE" if r["ceq_ok"] >= 0.90 else "HOLD")
print(f"\nfrontier stage {stage}: ceq_ok = {r['ceq_ok']:.3f} "
f"vs bar 0.90 -> {verdict}")
print(f" mean CE(q||p) = {r['ceq']:.4f}, floor H(q) = {r['h_q']:.4f}, "
f"ceiling log|S| = {r['log_m']:.4f}")
if args.out_json:
with open(args.out_json, "w") as f:
json.dump({"label": tag, "step": step, "eval_stage": stage,
"eval_epochs": args.eval_epochs, "per_stage": rows}, f,
indent=1)
print(f"wrote {args.out_json}")
if __name__ == "__main__":
main()
|