ONNX Exporter

Export PyTorch models to ONNX format with optional INT8 quantization

ONNX Export

Export PyTorch models to ONNX format for deployment. Supports: - Basic ONNX export with graph optimization - Dynamic INT8 quantization (no calibration needed) - Static INT8 quantization (with calibration data) - Faithful Q/DQ export of pt2e-quantized models - Output verification against original model

Export Function

export_onnx targets float models: it runs the ONNX graph optimizer, whose fusion passes remove the Q/DQ pairs of an already-quantized graph. To export a model quantized with Quantizer(backend='pt2e'), use export_qdq instead.

Inference Wrapper

Verification

Quantized (QDQ) Export

A model quantized with Quantizer(backend='pt2e') carries its quantization in the graph, as pairs of quantize / dequantize operators around each tensor. export_qdq translates those operators to the ONNX QuantizeLinear / DequantizeLinear pair and keeps them intact, so the ONNX file carries the same scales, zero-points and quantization axes as the PyTorch model — that is what downstream runtimes read to build their INT8 kernels (they remain free to fuse and reorder the operators around them). That translation table is only consulted by torch versions whose exporter does not already translate the quantized_decomposed operators itself; where it does, its own translation is the one that applies.

Two helpers come with it: qdq_stats reports what the exported file actually contains, and verify_qdq measures how often the ONNX graph and the PyTorch model agree on the predicted class.

Not every precision has a QDQ form: INT4 weights have no Q/DQ pair in opset 18, and a weight-only or dynamically-quantized model has no static activation scales to write. export_qdq reads the precision Quantizer recorded on the model and refuses those cases with a message naming the precision spec and the backend that produced it, rather than writing a graph that describes something else.

Everything export_qdq checks, it checks on the file it produced, not on the request that produced it:

  • The opset the file declares is the one that was asked for. torch.onnx exports at the exporter’s own opset and converts afterwards; when the ONNX version converter cannot rewrite an operator it keeps the original opset and only logs the failure. A pt2e graph runs into exactly one such operator: ReduceMean, whose axes moved from an attribute to an input in opset 18 — a move the converter cannot undo, so a graph that is opset 17 in every other respect stays at 18, and a parser limited to opset 17 rejects it. export_qdq moves the axes back itself, drops the constant that fed them, and then checks the file it wrote: onnx.checker reads it at the opset it now declares, and ONNX Runtime builds it and runs it — output-for-output identical to the graph as produced, or the rewrite is not kept. (Checking a lowered graph means running it, so that path needs onnxruntime.) What the rewrite cannot reach is refused with a message naming it — a ReduceMean whose axes are computed while the graph runs, or an operator that genuinely needs the newer opset — and nothing is kept on disk: not the graph, and not the external-data file it may reference.
  • Every Conv carries its kernel_shape. ONNX makes that attribute optional because a consumer can read the spatial dims off the weight tensor, and the dynamo exporter leaves it out. In a QDQ graph the weight is the output of a DequantizeLinear rather than an initializer, and an ONNX parser that only reads initializer shapes — TensorRT’s among them — then has nothing to read and rejects the graph. export_qdq resolves each Conv’s weight back through the Q/DQ chain and writes the attribute. This is spec-legal and changes no weight, no scale and no topology: strip the attributes back off and ONNX Runtime returns byte-identical logits. Only Conv is patched; ConvTranspose and the pooling operators are left exactly as the exporter wrote them.
  • activation_dtype='uint8' moves the activation pairs, and only those. Every QuantizeLinear and the DequantizeLinear nodes reading it take a uint8 zero-point of 128 with the scale they already had; the weight pairs keep their int8 zero-point of 0. The rewritten graph is written beside the produced one, read back node by node, checked by onnx.checker, and run — unoptimized against the produced graph, where the two must return byte-identical outputs, and once more at ONNX Runtime’s default optimization level, which it has to load and run. Anything else is refused and nothing is kept, not even the file as it was produced. The default 'int8' file is not touched, and a consumer that only accepts a zero-point of 0, TensorRT among them, wants that default.

Three details of that last rewrite. The comparison runs both graphs with the ONNX Runtime optimizations disabled: a fused execution requantizes in the runtime’s own arithmetic rather than through the Q/DQ pair the file carries. Byte-identity is what it asks for because the produced graph saturates to the full int8 range, so saturate(round(x/s) + 128, 0, 255) == saturate(round(x/s), -128, 127) + 128 — the same integers, read unsigned. And the produced graph types its QuantizeLinear outputs INT8 in value_info: the rewrite deletes those entries rather than retyping them, leaving the type where a consumer infers it anyway, from the node.

Usage Examples

from fasterai.export.all import export_onnx, ONNXModel, verify_onnx

# Basic export
path = export_onnx(model, sample, "model.onnx")

# With quantization
path = export_onnx(model, sample, "model.onnx", quantize=True)

# Inference
onnx_model = ONNXModel("model.onnx")
output = onnx_model(input_tensor)

# Verify
assert verify_onnx(model, "model.onnx", sample)

Exporting a pt2e-quantized model

from fasterai.quantize.quantizer import Quantizer
from fasterai.export.all import export_qdq, qdq_stats, verify_qdq

quantized_model = Quantizer(backend='pt2e').quantize(model, calibration_dl=dls.valid)
path = export_qdq(quantized_model, sample, "model_qdq.onnx")

# What did the exporter actually write?
print(qdq_stats(path).as_dict())
# for a ResNet-18: {'n_quantize': 33, 'n_dequantize': 54, 'n_per_channel': 21,
#                   'n_nonzero_zero_point': 0, 'n_unquantized_conv_add': 0, 'n_uint8': 0}
# (ResNet-18, torch 2.9.1, opset 18 — the Q/DQ counts track the exporter version; only n_per_channel,
#  n_nonzero_zero_point, n_unquantized_conv_add and n_uint8 are pinned by tests)

# A model quantized with Quantizer(backend='pt2e', qdq_placement='skip_conv_add') writes the same
# graph minus the pairs on the residual conv->add edges: on this ResNet-18, n_quantize 33 -> 25,
# n_dequantize 54 -> 46, and n_unquantized_conv_add 0 -> 8. That last count is also a POST-CONDITION:
# a spec asking for that placement whose file carries none of those edges is refused, and no file kept.

# For a parser that stops at opset 17: the graph is rewritten, checked and run before it comes back
path17 = export_qdq(quantized_model, sample, "model_qdq17.onnx", opset_version=17)
# Expect a logged traceback first: the ONNX version converter reports its own failure (torch.onnx runs it
# before export_qdq's rewrite ever sees the graph). It is not the outcome — the returned path is. On the
# day export_qdq refuses too, that traceback is the first diagnostic to read.

# For a consumer that reads unsigned activations (the weights stay int8, per channel):
path_u8 = export_qdq(quantized_model, sample, "model_qdq_u8.onnx", activation_dtype='uint8')
stats = qdq_stats(path_u8)
# every activation node now carries an explicit uint8 zero-point of 128, and nothing else moved:
# stats.n_uint8 > 0, stats.n_nonzero_zero_point == stats.n_uint8, stats.n_per_channel unchanged.

# Same predictions as PyTorch? (`sample` holds 8 batches worth of inputs)
print(verify_qdq(quantized_model, path, samples, n_batches=8))
# Agreement is only informative when the reference predictions vary — verify_qdq warns when they do not.

source

export_onnx

def export_onnx(
    model:nn.Module, sample:torch.Tensor, # Example input for tracing (with batch dim)
    output_path:str | Path, # Output .onnx file path
    *, opset_version:int=17, quantize:bool=False, # Apply INT8 quantization after export
    quantize_mode:str='dynamic', # "dynamic" (no calibration) or "static"
    calibration_data:Iterable | None=None, # DataLoader for static quantization
    optimize:bool=True, # Run ONNX graph optimizer
    dynamic_batch:bool=True, # Allow variable batch size at runtime
    input_names:list[str] | None=None, output_names:list[str] | None=None
)->Path:

Export a PyTorch model to ONNX format with optional quantization


source

ONNXModel

def ONNXModel(
    path:str | Path, device:str='cpu'
):

Wrapper for ONNX Runtime inference with PyTorch-like interface


source

verify_onnx

def verify_onnx(
    model:nn.Module, onnx_path:str | Path, sample:torch.Tensor, rtol:float=0.001, atol:float=1e-05
)->bool:

Verify ONNX model outputs match PyTorch model within tolerance


source

export_qdq

def export_qdq(
    model:nn.Module, # Quantized model, typically from `Quantizer(backend='pt2e')`
    sample:torch.Tensor, # Example input (with batch dim); use the calibration batch size
    output_path:str | Path, # Output .onnx file path
    *, opset_version:int=18, # ONNX opset the file must declare; a lower one is rewritten and verified
    activation_dtype:str='int8', # 'int8' (zero-point 0) or 'uint8' (zero-point 128); weights stay int8
    dynamic_batch:bool=False, # Ask for a dynamic batch dimension (experimental)
    input_names:list[str] | None=None, output_names:list[str] | None=None
)->Path:

Export a quantized model to ONNX, keeping its Q/DQ node pairs intact — with activation_dtype='uint8' the activation pairs are written with the same scales and zero-point 128, and the weight pairs stay int8.


source

QDQStats

def QDQStats(
    n_quantize:int, n_dequantize:int, n_per_channel:int, n_nonzero_zero_point:int, n_unquantized_conv_add:int,
    n_uint8:int=0
)->None:

Quantization nodes found in an ONNX graph, and the operator edges that carry no pair


source

qdq_stats

def qdq_stats(
    onnx_path:str | Path
)->QDQStats:

Count the Q/DQ nodes of an ONNX graph and check their zero-points


source

verify_qdq

def verify_qdq(
    model:nn.Module, # Reference PyTorch model (the quantized module that was exported)
    onnx_path:str | Path, sample:torch.Tensor, # Test inputs, batch dimension first
    n_batches:int=1, # Split `sample` into this many EQUAL batches
)->float:

Fraction of inputs on which the ONNX graph and the PyTorch model predict the same class


See Also

Tests live in nbs/tests/test_onnx_exporter.ipynb.