| """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"}} |
| |
| 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)) |
|
|
