Table of Contents
What ONNX Is and When to Use It
ONNX, the Open Neural Network Exchange format, is a standardized way to represent a trained model as a computation graph plus weights. In practice, exporting to ONNX is useful when you want to run a PyTorch model outside of PyTorch, for example in ONNX Runtime, in a different programming language, or in an environment where installing PyTorch is undesirable. ONNX is also commonly used as an interchange format on the way to other deployment targets.
An ONNX export is primarily an inference artifact. You typically export a model after training, and you expect the exported graph to run forward passes efficiently and consistently. Training behaviors such as dropout and batch normalization updates should be disabled at export time by putting the model in evaluation mode.
Always call model.eval() before exporting to ONNX. Otherwise, layers like Dropout and BatchNorm can change behavior and produce incorrect inference results.
Mental Model of the Export Process
Exporting to ONNX converts your model’s forward computation into a static graph of ONNX operators. PyTorch can do this in two broad ways: by tracing the execution with example inputs, or by capturing the Python level structure more directly. For beginners, the important takeaway is that export is guided by sample inputs, and the exporter must be able to express every operation in your forward pass using supported ONNX ops.
Because ONNX is a graph format, certain Python control flow patterns can be problematic, especially those that change computation based on data values. Also, operations that depend on Python side effects, custom C++ extensions, or unsupported PyTorch ops may fail to export or may export in a way that is not portable.
If an operation in forward() is not supported by ONNX, export can fail, or worse, succeed but produce a model that behaves differently. Always validate numerically after export.
A Minimal Export Example
A typical export starts with a trained torch.nn.Module, an example input tensor with the right shape and dtype, and a target file path. The following pattern is representative:
import torch
model.eval()
dummy_x = torch.randn(1, 3, 224, 224) # example shape for an image model
torch.onnx.export(
model,
dummy_x,
"model.onnx",
export_params=True,
opset_version=17,
do_constant_folding=True,
input_names=["input"],
output_names=["logits"],
dynamic_axes={"input": {0: "batch"}, "logits": {0: "batch"}},
)
The example input is not “training data”, it is a shape and type template that drives graph construction. The opset_version selects the ONNX operator set version to target. Newer opsets usually provide better coverage, but your runtime environment must support the chosen opset.
Dynamic Shapes and Batch Size
By default, many exports bake in the exact shapes seen in the example input, which can be inconvenient if you want variable batch sizes at inference time. The dynamic_axes argument allows you to mark particular dimensions as dynamic, most commonly the batch dimension. This is a practical default for deployment because it allows inference on any batch size without re exporting.
Be cautious when making non batch dimensions dynamic, such as sequence length for NLP or image height and width for vision. Some models can support this cleanly, but others rely on fixed shapes internally, and even if export succeeds, performance and compatibility can vary across runtimes.
Marking axes as dynamic does not guarantee the model is truly shape agnostic. It only tells the exporter to avoid hard coding those dimensions. Test with multiple real input shapes.
Common Export Pitfalls
A frequent source of failure is device and dtype mismatch during export. Export should generally be done on CPU unless you have a specific reason to export from GPU. Another common issue is exporting a model that returns non tensor objects, such as Python lists of tensors with varying lengths, dictionaries with non tensor values, or custom classes. ONNX expects outputs that can be represented as tensors, or structured outputs that the exporter explicitly supports.
Models that include randomness in the forward pass, or that behave differently depending on global flags, can also cause confusing results. Ensuring deterministic inference behavior, using model.eval(), and avoiding random ops in forward are key.
Verifying an ONNX Export
You should treat the ONNX file as a separate artifact that must be validated. A basic verification workflow is to run the PyTorch model and the ONNX model on the same input and compare outputs. Small floating point differences are normal, but large discrepancies indicate a broken export or an unsupported operation.
A typical numerical check uses ONNX Runtime:
import numpy as np
import onnxruntime as ort
import torch
model.eval()
x = torch.randn(4, 3, 224, 224)
with torch.no_grad():
torch_out = model(x).cpu().numpy()
sess = ort.InferenceSession("model.onnx", providers=["CPUExecutionProvider"])
onnx_out = sess.run(None, {"input": x.cpu().numpy()})[0]
max_diff = np.max(np.abs(torch_out - onnx_out))
print("max abs diff:", max_diff)The appropriate tolerance depends on the model and dtype. If you use mixed precision during training, export and verification should be done carefully in full precision unless you explicitly target reduced precision inference.
Do not assume export is correct because it produces a file. Always run a direct output comparison on representative inputs before deploying.
Operator Sets, Compatibility, and Portability
ONNX portability depends on three pieces aligning: the opset you export to, the operators your model uses, and the runtime you plan to run on. If you target ONNX Runtime, that is usually the most forgiving path. If you target another runtime, you may need to choose an opset and modeling patterns that the runtime supports well.
If export fails, the most practical approach is to simplify the forward pass, replace unsupported ops with supported equivalents, or adjust the model to produce simpler outputs. In some cases, adding small refactors such as avoiding in place operations, removing Python side logic from forward(), or making shapes explicit can make export succeed.
What ONNX Export Is Not
ONNX export is not the same as saving a PyTorch checkpoint, and it is not intended for resuming training in PyTorch. It also does not automatically include preprocessing and tokenization logic unless you explicitly build those steps into the exported graph, which is often nontrivial. For deployment readiness, you usually package preprocessing separately in application code, or you design a model wrapper that includes the required tensor level preprocessing.
ONNX is best thought of as a portable inference graph and weights, designed to be executed by a compatible runtime.