ONNX Runtime as a Portability Layer: Export, Verify, Deploy

ONNX gives you one model artefact that runs on CPU, GPU and several accelerators through a single API. What it does not give you is any guarantee that the exported graph computes what the original did, or that the provider you asked for is the one that ran. Both are checkable in about thirty lines, and both should be in your build.

What ONNX buys and what it does not

The bargain is: define the model once in a training framework, export it to a graph format with a versioned operator set, and let a runtime decide how to execute that graph on whatever hardware is present. In practice it removes a large amount of per-platform work — one artefact for a Windows desktop, a Linux server and an Android phone is a real saving.

The parts it does not solve, and which surprise people:

• Export is lossy in ways that pass silently. Tracing freezes control flow. Unsupported operators are decomposed into approximations. Numerics drift. Nothing errors.
• Not every provider supports every operator. The runtime partitions the graph and runs the unsupported parts elsewhere, so “I enabled the GPU provider” and “my model ran on the GPU” are different claims.
• Opset version is a compatibility surface. A newer opset can express your model more directly but may not be supported by an older runtime on a device you still ship to.

Exporting from PyTorch

import torch

model.eval()
example = torch.randn(1, 3, 224, 224)

torch.onnx.export(
    model,
    example,
    "model.onnx",
    opset_version=17,
    input_names=["input"],
    output_names=["logits"],
    dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
    do_constant_folding=True,
)

dynamic_axes is the argument to get right first. Without it, every dimension is frozen at the example’s size and a different-sized input fails at inference. With it applied too liberally — marking spatial dimensions dynamic when you only ever use three sizes — you deny the runtime the shape information it needs to optimise. Declare exactly the axes that genuinely vary.

Newer PyTorch releases also offer an exporter built on the compiler stack rather than on tracing, which handles some dynamic control flow that tracing cannot. Whether it is available and what the flag is called depends on your version; check the exporter documentation for the PyTorch you have pinned rather than copying a flag from a blog post.

Then check the file is structurally valid before going further:

import onnx

m = onnx.load("model.onnx")
onnx.checker.check_model(m)
print(f"opset: {m.opset_import[0].version}")
print(f"nodes: {len(m.graph.node)}")

A node count far larger than the number of layers in your model is a useful early warning: it usually means an operator was decomposed into a long sequence of primitives, which is both slower and more likely to be placed on CPU.

Verifying the numerics

This is the step that separates a working export from one that looks like it works. Run both implementations on the same inputs and compare with an explicit tolerance.

import numpy as np
import onnxruntime as ort
import torch

sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])

rng = np.random.default_rng(0)
max_abs = 0.0
mismatched_top1 = 0

for _ in range(64):
    x = rng.standard_normal((1, 3, 224, 224)).astype(np.float32)

    with torch.no_grad():
        ref = model(torch.from_numpy(x)).numpy()

    got = sess.run(["logits"], {"input": x})[0]

    max_abs = max(max_abs, float(np.abs(ref - got).max()))
    if ref.argmax(-1) != got.argmax(-1):
        mismatched_top1 += 1

print(f"max abs diff: {max_abs:.3e}")
print(f"top-1 mismatches: {mismatched_top1}/64")
assert max_abs < 1e-4, "export changed the numerics"
assert mismatched_top1 == 0, "export changed the decision"

Random inputs are deliberate: they cover the input space more evenly than a handful of real samples and they catch operators that behave correctly only near the data distribution. Add real samples on top — not instead.

Assert on both the numeric difference and the decision. A classification model can drift by 1e-3 in logits without changing a single prediction, which is fine, or drift by 1e-5 at exactly the decision boundary of your most common class, which is not. The second assertion is what catches that. Put both in continuous integration so that a framework upgrade cannot quietly change your model.

Execution providers are a preference list

import onnxruntime as ort

print(ort.get_available_providers())

sess = ort.InferenceSession(
    "model.onnx",
    providers=["CUDAExecutionProvider", "CPUExecutionProvider"],
)

print(sess.get_providers())   # what was actually registered

The list you pass is ordered by preference, and the runtime falls back quietly down it. Two habits follow. First, always call get_providers() after construction and log it — otherwise a missing driver becomes a mysterious tenfold slowdown rather than a line in a log. Second, always include CPUExecutionProvider last: it is the only one guaranteed to be there, and omitting it turns a degraded run into a crash.

Availability is a property of the build, not of ONNX Runtime in general. The package that ships CPU only and the package that ships a particular accelerator are different installs, and installing both in one environment is a reliable way to spend an afternoon.

Proving the graph actually ran where you asked

Registering a provider tells you it was available. It does not tell you how much of your graph it took. Turn on profiling and read the trace:

so = ort.SessionOptions()
so.enable_profiling = True

sess = ort.InferenceSession("model.onnx", so,
                            providers=["CUDAExecutionProvider",
                                       "CPUExecutionProvider"])

for _ in range(20):
    sess.run(None, {"input": x})

trace_path = sess.end_profiling()   # writes a JSON trace
print(trace_path)

The trace records per-node execution with the provider that ran it. Sum the time per provider. If a graph you believed was fully accelerated shows a third of its time on CPU nodes, you have found the partition boundary and the operator that caused it — and usually the fix is to change that one operator in the source model and re-export, not to fight the runtime.

Leave profiling off in production. It writes a trace file per session and the overhead is real.

Quantising an ONNX model

ONNX Runtime ships quantisation tooling that operates on the exported graph, which means you can compress without returning to the training framework. Two modes, and the choice is the usual one:

If your target is an NPU, you almost certainly need static: integer-only hardware cannot compute activation ranges on the fly. Give the calibrator a few hundred genuinely representative inputs — not the training set’s first batch, which is often sorted. Then re-run the verification from the section above against the quantised model, with a tolerance appropriate to int8 and a hard assertion on decisions.

The session options that change the numbers

A default session is not a tuned session, and on edge hardware the gap is large. Four settings account for most of it:

so = ort.SessionOptions()
so.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
so.intra_op_num_threads = 4      # threads within one operator
so.inter_op_num_threads = 1      # operators run concurrently
so.enable_mem_pattern = True     # pre-plan allocations for static shapes

sess = ort.InferenceSession("model.onnx", so,
                            providers=["CPUExecutionProvider"])

Save the optimised graph once rather than re-optimising on every launch: setting an optimised-model output path writes the transformed graph to disk, and loading that afterwards skips the work. Do note that the saved graph is specialised to the optimisations and providers available when it was written, so it belongs in a cache keyed on the runtime version rather than in your repository.

Export pitfalls worth knowing in advance

• Forgetting model.eval(). Exporting in training mode bakes dropout and batch-norm running-statistics updates into the graph. The model works and is subtly wrong. This is the most common export bug there is.
• Python control flow on tensor values. if x.sum() > 0: in a forward pass is resolved at trace time and frozen. If your model branches on data, it must be expressed with graph-level control flow or split into two models.
• Preprocessing left outside the graph. Normalisation constants, colour channel order and resize interpolation are part of the model as far as correctness is concerned. Either export them into the graph or write them down as part of the artefact’s contract. A mismatch here produces a model that is merely mediocre rather than broken, which is why it survives review.
• Opset chosen by default. Pin it explicitly. An upgrade of your training framework should not silently change the operator set your device runtime has to support.
• Assuming the artefact is portable to a runtime you have not tested. ONNX is a portable format, not a portability guarantee. Run the verification script on every target runtime you ship to, including the mobile one, and treat a new target as a new integration.