| |
| """Export NadirNet to ONNX: the page graph (palette + page context) and the element graph (treatment + role for any |
| set of page images). Calibrated 8-bit and 4-bit variants of the page graph come from quantize.py.""" |
| import argparse, json, os, pickle, time |
| import numpy as np |
| import torch |
| import torch.nn as nn |
|
|
| from model import NadirNet, ElemNet, FI, TREAT, IMG_ROLES, N_CTX |
| from train import collate, prep, element_logits |
|
|
| NAMES_IN = ["tok_f", "tok_role", "tok_lab", "img_f", "img_thumb", "page_f", "src", "rel", "mask"] |
| NAMES_OUT = ["lab", "alpha_logit", "ctx"] |
| ELEM_IN = ["e_f", "e_thumb", "under_light", "under_dark", "ctx"] |
| ELEM_OUT = ["treat_logits", "role_logits"] |
|
|
|
|
| class Wrapped(nn.Module): |
| def __init__(self, m): |
| super().__init__() |
| self.m = m |
|
|
| def forward(self, tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask): |
| return self.m(tok_f, tok_role, tok_lab, img_f, img_thumb, page_f, src, rel, mask) |
|
|
|
|
| def feeds_for(s, pad=None): |
| b = collate([s], "cpu") |
| if pad: |
| N0, M0 = b["tok_f"].shape[1], b["img_f"].shape[1] |
| N, M = pad |
| if N0 <= N and M0 <= M: |
| T0, T = 1 + N0 + M0, 1 + N + M |
| def padto(x, dims, val=0): |
| shape = list(x.shape) |
| for d, sz in dims.items(): |
| shape[d] = sz |
| y = torch.full(shape, val, dtype=x.dtype) |
| y[tuple(slice(0, x.shape[i]) for i in range(x.ndim))] = x |
| return y |
| idx = list(range(1 + N0)) + list(range(1 + N, 1 + N + M0)) |
| rel = torch.zeros(1, T, T, b["rel"].shape[-1]); mask = torch.zeros(1, T, dtype=torch.bool) |
| ii = torch.tensor(idx) |
| rel[0][ii[:, None], ii[None, :]] = b["rel"][0] |
| mask[0, ii] = b["mask"][0] |
| b = dict(b, tok_f=padto(b["tok_f"], {1: N}), role=padto(b["role"], {1: N}), lab=padto(b["lab"], {1: N}), |
| img_f=padto(b["img_f"], {1: M}), thumbs=padto(b["thumbs"], {1: M}), rel=rel, mask=mask) |
| M = b["img_f"].shape[1] |
| if M == 0: |
| T = b["mask"].shape[1] |
| b["img_f"] = torch.zeros(1, 1, FI); b["thumbs"] = torch.zeros(1, 1, 4, 32, 32) |
| rel = torch.zeros(1, T + 1, T + 1, b["rel"].shape[-1]); rel[:, :T, :T] = b["rel"]; b["rel"] = rel |
| b["mask"] = torch.cat([b["mask"], torch.zeros(1, 1, dtype=torch.bool)], 1) |
| return {"tok_f": b["tok_f"].numpy(), "tok_role": b["role"].numpy(), "tok_lab": b["lab"].numpy(), "img_f": b["img_f"].numpy(), |
| "img_thumb": b["thumbs"].numpy(), "page_f": b["page"].numpy(), "src": b["src"].numpy(), "rel": b["rel"].numpy(), "mask": b["mask"].numpy()} |
|
|
|
|
| def elem_feeds(s, lab, ctx): |
| """Element-graph inputs for a page's labelled images, with the dark surface taken from the page graph's output.""" |
| b = collate([s], "cpu") |
| k = int(b["e_mask"][0].sum()) |
| ut = b["e_ut"][0, :k].numpy() |
| return {"e_f": b["e_f"][:, :k].numpy(), "e_thumb": b["e_thumb"][:, :k].numpy(), "under_light": b["e_ul"][:, :k].numpy(), |
| "under_dark": lab[:, ut].astype(np.float32), "ctx": ctx.astype(np.float32)}, b["e_soft"][0, :k].numpy() |
|
|
|
|
| def fidelity(sess, ref_sess, samples): |
| """Mean OKLab distance between this model's colours and fp32's (weighted like training), worst token, ctx drift.""" |
| de, worst, ctxd = [], 0.0, [] |
| for s in samples: |
| f = feeds_for(s) |
| a = sess.run(None, f); r = ref_sess.run(None, f) |
| n = s["F"].shape[0] |
| d = np.linalg.norm(a[0][0, :n] - r[0][0, :n], axis=-1) |
| w = s["w"][:n] |
| de.append(float((d * w).sum() / w.sum())); worst = max(worst, float(d.max())) |
| ctxd.append(float(np.abs(a[2] - r[2]).max())) |
| return {"dE_vs_fp32": float(np.mean(de)), "dE_max": worst, "ctx_max_abs": float(np.max(ctxd))} |
|
|
|
|
| def bench(sess, feeds, reps=3): |
| ts = [] |
| for f in feeds[:20]: |
| sess.run(None, f) |
| t = time.time() |
| for _ in range(reps): |
| sess.run(None, f) |
| ts.append((time.time() - t) / reps * 1000) |
| return float(np.median(ts)) |
|
|
|
|
| def main(): |
| ap = argparse.ArgumentParser() |
| ap.add_argument("--ckpt", required=True) |
| ap.add_argument("--data", required=True) |
| ap.add_argument("--out", required=True) |
| ap.add_argument("--only-float", action="store_true", help="kept for the job scripts (quantization lives in quantize.py)") |
| args = ap.parse_args() |
| os.makedirs(args.out, exist_ok=True) |
| import onnx, onnxruntime as ort |
| ck = torch.load(args.ckpt, map_location="cpu") |
| model = NadirNet(**{k: v for k, v in ck["cfg"].items()}); model.load_state_dict(ck["model"]); model.eval() |
| data = pickle.load(open(args.data, "rb")) |
| val = [prep(s) for s in data["val"]] |
| s0 = max(val, key=lambda s: len(s["role"])) |
| ex = feeds_for(s0) |
| par_samples = val[:8] |
| with torch.no_grad(): |
| ref_out = [model(*[torch.from_numpy(feeds_for(s)[k]) for k in NAMES_IN]) for s in par_samples] |
| fp32 = os.path.join(args.out, "nadir-fp32.onnx") |
| torch.onnx.export(Wrapped(model), tuple(torch.from_numpy(ex[k]) for k in NAMES_IN), fp32, input_names=NAMES_IN, output_names=NAMES_OUT, |
| dynamic_axes={"tok_f": {1: "N"}, "tok_role": {1: "N"}, "tok_lab": {1: "N"}, "img_f": {1: "M"}, "img_thumb": {1: "M"}, |
| "rel": {1: "T", 2: "T"}, "mask": {1: "T"}, "lab": {1: "N"}, "alpha_logit": {1: "N"}}, |
| opset_version=17, do_constant_folding=True, dynamo=False) |
| onnx.checker.check_model(fp32) |
| so = ort.SessionOptions(); so.intra_op_num_threads = 4 |
| ref = ort.InferenceSession(fp32, so, providers=["CPUExecutionProvider"]) |
| parity = 0.0 |
| for r, s in zip(ref_out, par_samples): |
| o = ref.run(None, feeds_for(s)) |
| parity = max(parity, float(np.abs(r[0].numpy() - o[0]).max()), float(np.abs(r[2].numpy() - o[2]).max())) |
| |
| es = [s for s in val if len(s["E_ut"])][:64] |
| ef = [] |
| for s in es: |
| o = ref.run(None, feeds_for(s)) |
| ef.append(elem_feeds(s, o[0], o[2])) |
| exe = max(ef, key=lambda x: x[0]["e_f"].shape[1])[0] if ef else {"e_f": np.zeros((1, 4, FI), np.float32), "e_thumb": np.zeros((1, 4, 4, 32, 32), np.float32), |
| "under_light": np.zeros((1, 4, 3), np.float32), "under_dark": np.zeros((1, 4, 3), np.float32), "ctx": np.zeros((1, N_CTX), np.float32)} |
| elem = os.path.join(args.out, "nadir-elem-fp32.onnx") |
| em = ElemNet(model).eval() |
| torch.onnx.export(em, tuple(torch.from_numpy(exe[k]) for k in ELEM_IN), elem, input_names=ELEM_IN, output_names=ELEM_OUT, |
| dynamic_axes={"e_f": {1: "K"}, "e_thumb": {1: "K"}, "under_light": {1: "K"}, "under_dark": {1: "K"}, "treat_logits": {1: "K"}, "role_logits": {1: "K"}}, |
| opset_version=17, do_constant_folding=True, dynamo=False) |
| onnx.checker.check_model(elem) |
| esess = ort.InferenceSession(elem, so, providers=["CPUExecutionProvider"]) |
| eparity, agree, ok, n = 0.0, 0, 0, 0 |
| T = float(ck.get("temperature", 1.0)) |
| for f, soft in ef: |
| with torch.no_grad(): |
| tt, rr = em(*[torch.from_numpy(f[k]) for k in ELEM_IN]) |
| o = esess.run(None, f) |
| eparity = max(eparity, float(np.abs(tt.numpy() - o[0]).max()), float(np.abs(rr.numpy() - o[1]).max())) |
| pred = o[0][0].argmax(-1) |
| agree += int((pred == soft.argmax(-1)).sum()); ok += int((soft[np.arange(len(pred)), pred] > 0).sum()); n += len(pred) |
| report = {"params": sum(p.numel() for p in model.parameters()), |
| "elem_params": sum(p.numel() for k, p in model.named_parameters() if k.startswith(("cnn.", "img_in.", "elem_head."))), |
| "parity_max_abs": parity, "elem_parity_max_abs": eparity, "temperature": T, "treat": TREAT, "img_roles": IMG_ROLES, |
| "elem_val_sample": {"n": n, "acc": agree / max(1, n), "ok": ok / max(1, n)}, "variants": {}} |
| print("parity", parity, "elem parity", eparity, flush=True) |
| r = fidelity(ref, ref, val[:50]); r["size_mb"] = round(os.path.getsize(fp32) / 1e6, 2); r["cpu_ms_median"] = round(bench(ref, [feeds_for(s) for s in val[:20]]), 2) |
| report["variants"]["fp32"] = r |
| report["elem"] = {"size_mb": round(os.path.getsize(elem) / 1e6, 2), "cpu_ms_median": round(bench(esess, [f for f, _ in ef]), 2) if ef else None} |
| print(report, flush=True) |
| json.dump(report, open(os.path.join(args.out, "export_report.json"), "w"), indent=1) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
