File size: 4,900 Bytes
45507bb | 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 | """PyTorch checkpoint -> ONNX (+ parity check, + optional INT8 dynamic quantization).
The exported graph takes (state, question_ids, choice_ids) with dynamic batch,
number of choices and state size, and returns (probs, logits, embedding); the
calibration temperature is baked in. Everything the runtime needs to
pre-process inputs is stored in the model's metadata (key "minijev").
"""
from __future__ import annotations
import json
from pathlib import Path
import numpy as np
import torch
from minijev.inference.runtime import META_KEY, OnnxDecider, TorchDecider
from minijev.model import STATE_INPUTS, ExportWrapper
def export_onnx(checkpoint: str | Path, out: str | Path | None = None, opset: int = 18,
verify: bool = True, quantize: bool = True) -> dict:
import onnx
checkpoint = Path(checkpoint)
out = Path(out) if out else checkpoint.with_suffix(".onnx")
torch_dec = TorchDecider(checkpoint)
model, cfg, task = torch_dec.model, torch_dec.cfg, torch_dec.task
state_key = STATE_INPUTS[task.state_encoder]
wrapper = ExportWrapper(model, state_key).eval()
states = task.example_states(4)
question = task.default_question
choices = list(task.default_choices)
x = torch_dec.encode(states)
q, c = torch_dec._tokens(question, tuple(choices))
B = len(states)
args = (torch.as_tensor(x), torch.as_tensor(np.repeat(q, B, 0)), torch.as_tensor(np.repeat(c, B, 0)))
state_axes = {0: "batch", 2: "height", 3: "width"} if x.ndim == 4 else {0: "batch", 1: "length"}
dynamic_axes = {state_key: state_axes, "question_ids": {0: "batch"},
"choice_ids": {0: "batch", 1: "n_choices"},
"probs": {0: "batch", 1: "n_choices"}, "logits": {0: "batch", 1: "n_choices"},
"embedding": {0: "batch"}}
# TorchScript-based exporter: mature, handles SDPA with masks and dynamic axes well
torch.onnx.export(wrapper, args, str(out), input_names=[state_key, "question_ids", "choice_ids"],
output_names=["probs", "logits", "embedding"], dynamic_axes=dynamic_axes,
opset_version=opset, dynamo=False)
meta = {"task": task.name, "task_cfg": cfg["task"], "state_input": state_key,
"max_question_len": torch_dec.max_question_len, "max_choice_len": torch_dec.max_choice_len,
"temperature": model.temperature, "params": model.num_parameters(), "config_name": cfg.get("name"),
"default_question": question, "default_choices": choices}
m = onnx.load(str(out))
entry = m.metadata_props.add()
entry.key, entry.value = META_KEY, json.dumps(meta)
onnx.save(m, str(out))
report = {"onnx": str(out), "size_mb": round(out.stat().st_size / 2**20, 3), "meta": meta}
if verify:
report["parity_fp32"] = parity(torch_dec, OnnxDecider(out), task)
if quantize:
q_out = out.with_name(out.stem + ".int8.onnx")
quantize_int8(out, q_out)
report["onnx_int8"] = str(q_out)
report["size_int8_mb"] = round(q_out.stat().st_size / 2**20, 3)
if verify:
report["parity_int8"] = parity(torch_dec, OnnxDecider(q_out), task)
return report
def parity(ref, other, task, n: int = 256, seed: int = 1) -> dict:
"""Compare probabilities on real states, including 2- and 3-choice subsets."""
states = task.example_states(n, seed=seed)
choices = list(task.default_choices)
res = {}
for name, ch in (("all", choices), ("subset", choices[: max(2, len(choices) - 1)])):
p_ref = ref.probs(states, task.default_question, ch)
p_oth = other.probs(states, task.default_question, ch)
res[name] = {"max_abs_diff": float(np.abs(p_ref - p_oth).max()),
"argmax_agreement": float((p_ref.argmax(1) == p_oth.argmax(1)).mean())}
return res
def quantize_int8(src: str | Path, dst: str | Path) -> None:
"""Dynamic INT8 quantization of MatMul/Gemm weights (activations stay float).
Conv layers are left in fp32: dynamic ConvInteger kernels are slow on most
CPUs and the stem is a small share of the compute.
"""
from onnxruntime.quantization import QuantType, quantize_dynamic
from onnxruntime.quantization.shape_inference import quant_pre_process
pre = Path(str(dst) + ".pre.onnx")
try:
quant_pre_process(str(src), str(pre), skip_symbolic_shape=True)
source = pre
except Exception:
source = Path(src)
quantize_dynamic(str(source), str(dst), weight_type=QuantType.QInt8, op_types_to_quantize=["MatMul", "Gemm"])
if pre.exists():
pre.unlink()
import onnx
m = onnx.load(str(dst))
src_meta = {p.key: p.value for p in onnx.load(str(src)).metadata_props}
del m.metadata_props[:]
for k, v in src_meta.items():
e = m.metadata_props.add()
e.key, e.value = k, v
onnx.save(m, str(dst))
|