ONNX

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

Common Mistakes

• Exporting to ONNX without setting dynamic_axes for 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?

Related DL Notes

• TorchScript
• Model Optimization (Deployment)
• GPU Deployment