File size: 8,680 Bytes
bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 864f427 bde89e0 | 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 | #!/usr/bin/env python3
"""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: # fixed shapes for calibration (quantizers stack activations across samples)
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: # the runtime always sends at least one (masked) image slot
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()))
# element graph
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()
|