ByteAtlas

Exporting a PyTorch model to ONNX can look finished when a .onnx file appears on disk. That is only the first boundary. The harder question is whether the graph accepts the shapes and data types your service will actually send, produces numerically acceptable outputs, and runs on the execution provider you intended.

In model-integration work, I treat export as a contract migration rather than a file conversion. PyTorch owns one execution contract; ONNX and ONNX Runtime own another. Most surprises live between them: a batch dimension frozen to one, preprocessing that never entered the graph, a CUDA session quietly falling back to CPU, or output differences that only appear on realistic inputs.

The workflow I use

• Put the module in inference mode and choose representative inputs.
• Export with the torch.export-based ONNX path.
• Inspect the saved model’s inputs, outputs, and symbolic dimensions.
• Create an ONNX Runtime session with an explicit provider order.
• Compare PyTorch and ONNX Runtime outputs on more than one input.
• Only then measure latency or move the model into a service.

This article uses a small image classifier so the deployment mechanics stay visible. The same checks apply to larger vision, speech, and transformer models, although those models often add multiple inputs, caches, or custom operators.

Install the export and runtime packages

python -m venv .venv
source .venv/bin/activate
python -m pip install --upgrade pip
python -m pip install torch onnx onnxruntime onnxscript numpy

Use onnxruntime-gpu instead of onnxruntime for a supported NVIDIA CUDA environment. Do not install both and assume Python will select the right one. After installation, onnxruntime.get_available_providers() tells you what the installed runtime can actually use.

I intentionally avoid pinning versions in the command above. The exporter changes with PyTorch releases, while CUDA-enabled ONNX Runtime packages have their own compatibility requirements. In a real project, resolve a working set, record it in the lock file, and promote that exact set through CI.

Start with an exportable inference module

import torch
from torch import nn

class TinyClassifier(nn.Module):
    def __init__(self, num_classes: int = 10):
        super().__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 16, kernel_size=3, padding=1),
            nn.ReLU(),
            nn.AdaptiveAvgPool2d((1, 1)),
        )
        self.classifier = nn.Linear(16, num_classes)

    def forward(self, images: torch.Tensor) -> torch.Tensor:
        features = self.features(images)
        return self.classifier(features.flatten(1))

model = TinyClassifier()
model.eval()
example_images = torch.randn(2, 3, 224, 224, dtype=torch.float32)

model.eval() is easy to overlook because export may still complete without it. The problem is semantic: dropout and batch normalization should behave as they will during inference. Also keep the example input representative. Its rank, dtype, and non-dynamic dimensions become part of the exported contract.

If production applies resize, channel reordering, normalization, tokenization, or padding outside the module, the ONNX graph will not magically contain those operations. Document that preprocessing beside the model or deliberately wrap supported preprocessing into the exported module.

Export with deliberate dynamic shapes

batch = torch.export.Dim("batch", min=1, max=32)
height = torch.export.Dim("height", min=32)
width = torch.export.Dim("width", min=32)

onnx_program = torch.onnx.export(
    model,
    (example_images,),
    input_names=["images"],
    output_names=["logits"],
    dynamic_shapes={
        "images": {0: batch, 2: height, 3: width}
    },
    dynamo=True,
    verify=True,
    report=True,
)
onnx_program.save("tiny_classifier.onnx")

The current PyTorch exporter uses the torch.export-based path when dynamo=True. With that path, dynamic_shapes is the relevant option; older examples built around the legacy exporter commonly use dynamic_axes. Mixing the styles is a frequent source of confusion when code has been copied across PyTorch generations.

I do not mark every axis dynamic by reflex. A model trained for three-channel images should usually keep the channel dimension fixed. Batch, height, and width are dynamic here because this module can support them. If later code reshapes with a hard-coded spatial size, an export constraint or runtime failure is useful information rather than something to suppress.

verify=True asks the exporter to verify the program with ONNX Runtime, while report=True produces an export report when capture or translation fails. I still perform an independent comparison because exporter verification is not a substitute for testing the actual input distribution.

Check the ONNX file before creating a session

import onnx

onnx_model = onnx.load("tiny_classifier.onnx")
onnx.checker.check_model(onnx_model)

for value in onnx_model.graph.input:
    print("input:", value.name)
for value in onnx_model.graph.output:
    print("output:", value.name)

The ONNX checker validates the model against the ONNX specification. It does not prove numerical equivalence and it does not prove that a particular execution provider supports every node. Treat it as a structural gate, not the finish line.

Run the graph with ONNX Runtime

import onnxruntime as ort

session = ort.InferenceSession(
    "tiny_classifier.onnx",
    providers=["CPUExecutionProvider"],
)

print("active providers:", session.get_providers())
for item in session.get_inputs():
    print("input:", item.name, item.shape, item.type)
for item in session.get_outputs():
    print("output:", item.name, item.shape, item.type)

ort_inputs = {
    session.get_inputs()[0].name:
        example_images.detach().cpu().numpy()
}
ort_logits = session.run(None, ort_inputs)[0]

Using session metadata instead of repeating a guessed input name has saved more time than any clever exporter flag. Names can change when wrappers are introduced, and multi-input models make positional assumptions brittle. Metadata also exposes the expected tensor type; a NumPy float64 array will not satisfy an input declared as tensor(float).

For CUDA, put CUDAExecutionProvider before CPUExecutionProvider and check it is present in ort.get_available_providers(). CPU fallback may be useful, but it can also hide a packaging problem or unsupported CUDA subgraph. If GPU execution is a requirement, fail startup when the provider is unavailable rather than publishing misleading performance numbers.

Compare outputs, not just class IDs

import numpy as np

with torch.inference_mode():
    torch_logits = model(example_images).cpu().numpy()

np.testing.assert_allclose(
    torch_logits,
    ort_logits,
    rtol=1e-4,
    atol=1e-5,
)
print("PyTorch and ONNX Runtime logits are close.")

Comparing only argmax can miss meaningful drift in logits. A class label may remain the same while margins or downstream thresholds move. The tolerances above are an example starting point, not a universal acceptance rule. Precision, model depth, operator implementations, and the business decision built on the output all matter.

Repeat the comparison with several batch sizes, boundary shapes, empty or padded regions where valid, and samples from the same preprocessing path used by the service. For a detection model, compare decoded boxes and scores after matching detections. For an embedding model, check the downstream distance or ranking metric, not only element-wise error.

Failures that tend to surface late

The model exported with a fixed batch size

If batch one works and batch eight fails, inspect the session input shape. Dynamic dimensions must be declared during export, and the graph itself must avoid shape assumptions that contradict them. Editing only the ONNX input metadata does not make the internal graph dynamic.

The runtime rejects the input dtype

NumPy defaults and preprocessing libraries can quietly produce float64 or integer arrays. Cast intentionally and compare that choice with session.get_inputs(). Token IDs may need int64 even when another part of the application uses 32-bit integers.

An operator cannot be translated or executed

Separate exporter support from runtime-provider support. An operator may fail while PyTorch lowers the program to ONNX, or it may exist in ONNX but lack a kernel on the selected provider. The export report helps with the first case. ONNX Runtime logs and provider documentation help with the second. A custom operator is possible, but it creates a deployment dependency that should be chosen deliberately.

Python control flow depends on tensor values

Branches and loops that depend on tensor data need graph-compatible control flow. Code that is natural in eager PyTorch may be specialized to the example input or rejected during export. Reduce the module to the smallest failing forward path before trying random opset versions.

The graph is right, but the service is wrong

Channel order, normalization constants, tokenizer settings, padding side, and output decoding often live outside the graph. A clean PyTorch-versus-ORT comparison that feeds the same already-processed tensor will not catch a mismatch introduced earlier. Keep at least one end-to-end test that begins with the application’s raw input.

Performance comes after correctness

Create the session once and reuse it. Warm up before timing. Synchronize the device when measuring asynchronous GPU work, keep preprocessing and transfer boundaries explicit, and measure the batch sizes your application will use. Otherwise, an apparent improvement may simply be a different timing boundary.

ONNX Runtime applies graph optimizations and partitions work across available execution providers. That can improve deployment performance, but the runtime documentation is explicit that users remain responsible for validating accuracy, performance, and suitability. A portable graph is not a performance guarantee.

Deployment handoff checklist

• Record PyTorch, ONNX, ONNX Runtime, and provider package versions.
• Store input and output names, dtypes, shapes, and dynamic constraints.
• Version preprocessing and postprocessing with the model.
• Compare outputs on representative and boundary inputs.
• Confirm active execution providers on the target host.
• Decide whether CPU fallback is allowed or should fail fast.
• Measure after warmup using the real request shape distribution.

Conclusion

A reliable PyTorch-to-ONNX workflow is not export() followed by InferenceSession(). It is a sequence of contracts: inference behavior, input shapes, tensor types, preprocessing, numerical tolerance, and provider selection. Make each contract visible and failures become much easier to isolate.

For API details, use the official PyTorch ONNX exporter documentation, the PyTorch ONNX tutorial, and the ONNX Runtime Python guide. The same model contract and validation principles carry into a C++ deployment even though session construction and tensor ownership change.