ONNX (Open Neural Network Exchange) is a framework-agnostic model format — a way to export a model trained in PyTorch (or TensorFlow, or other frameworks) into a common representation that many different runtimes and hardware backends can execute.
Why a Framework-Agnostic Format Matters
Without ONNX, a model trained in PyTorch is generally tied to PyTorch-compatible deployment environments. ONNX defines a standard computational graph representation that many specialized inference runtimes (ONNX Runtime, TensorRT, and others) can consume directly, often with significant inference speed optimizations that a general-purpose training framework doesn't apply by default — genuinely valuable when raw inference speed matters in production.
Code — Exporting a PyTorch Model to ONNX
import torch
model.eval()
dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
model,
dummy_input,
"model.onnx",
input_names=["input"],
output_names=["output"],
dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
# dynamic_axes lets the exported model accept variable batch sizes at inference,
# rather than being locked to the exact batch size used during export
)Code — Running Inference with ONNX Runtime
import onnxruntime as ort
import numpy as np
session = ort.InferenceSession("model.onnx")
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
outputs = session.run(None, {"input": input_data})
print(outputs[0].shape)Notice that this inference code has no dependency on PyTorch at all — the ONNX Runtime session loads and executes the exported graph independently, which is exactly the portability benefit ONNX provides.
ONNX vs TorchScript — When to Use Which
| TorchScript | ONNX | |
|---|---|---|
| Ecosystem | PyTorch-specific (LibTorch for C++ deployment) | Framework-agnostic — usable across many runtimes and hardware backends |
| Control flow support | Full support via scripting | More limited — complex dynamic control flow can be harder to export correctly |
| Typical use case | Deploying specifically within a PyTorch/LibTorch-based system | Cross-framework deployment, or targeting specialized inference hardware/runtimes |
Common Mistakes
- Exporting to ONNX without setting
dynamic_axesfor the batch dimension, then discovering the exported model only accepts the exact batch size used during export. - Assuming an ONNX export will automatically support every possible PyTorch operation — some custom or exotic operations may not have direct ONNX equivalents and require special handling or workarounds.
Interview Relevance
Q: "Why might a team choose to export a PyTorch model to ONNX before deploying it, rather than deploying the native PyTorch model directly?" ONNX provides framework-agnostic portability — the exported model can run on specialized, often significantly faster inference runtimes (ONNX Runtime, TensorRT) and various hardware backends without requiring a PyTorch dependency in the serving environment. This can deliver meaningful inference speed and deployment flexibility benefits, particularly valuable when raw latency or cross-platform compatibility matters in production.
Practice Question
Why is setting dynamic_axes during ONNX export important for a model that will serve requests with varying batch sizes in production?
