PyTorch ONNX Export

ONNX (Open Neural Network Exchange) is an open format for representing machine learning models. Exporting a PyTorch model to ONNX lets you run it in environments that don't support PyTorch at all — TensorRT for NVIDIA GPUs, OpenVINO for Intel hardware, ONNX Runtime for cross-platform CPU/GPU inference, and CoreML for Apple devices.

Why Export to ONNX

PyTorch model:
  ✓ Best for research and training
  ✗ Requires PyTorch runtime
  ✗ Not supported on many deployment targets

ONNX model:
  ✓ Framework-independent file format
  ✓ Runs on ONNX Runtime, TensorRT, OpenVINO, CoreML
  ✓ Often faster inference than native PyTorch via specialized runtimes
  ✓ Enables deployment on hardware that has no Python or PyTorch

Use ONNX when your inference environment doesn't support PyTorch
or when a specialized runtime offers better performance.

The ONNX Export Flow

PyTorch Model
     │
     ▼  torch.onnx.export(model, dummy_input, 'model.onnx')
ONNX Graph (.onnx file)
     │
     ├──▶ ONNX Runtime  (CPU / GPU, cross-platform)
     ├──▶ TensorRT      (NVIDIA GPU, ultra-fast)
     ├──▶ OpenVINO      (Intel CPU / VPU)
     └──▶ CoreML        (Apple iPhone / Mac)

Exporting a Model to ONNX

import torch
import torch.nn as nn

class Net(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(10, 32)
        self.fc2 = nn.Linear(32, 3)

    def forward(self, x):
        return self.fc2(torch.relu(self.fc1(x)))

model = Net()
model.eval()   # MUST set eval mode before export

# Create a dummy input that matches the expected input shape
dummy_input = torch.rand(1, 10)   # batch=1, features=10

torch.onnx.export(
    model,
    dummy_input,
    'net.onnx',
    input_names=['input'],          # name the input node
    output_names=['output'],        # name the output node
    dynamic_axes={
        'input':  {0: 'batch_size'},  # make batch dimension dynamic
        'output': {0: 'batch_size'}
    },
    opset_version=17                # ONNX opset version
)

print("Model exported to net.onnx")

Verifying the ONNX File

import onnx

model_onnx = onnx.load('net.onnx')
onnx.checker.check_model(model_onnx)   # raises exception if invalid
print("ONNX model is valid")
print(onnx.helper.printable_graph(model_onnx.graph))

Running Inference with ONNX Runtime

import onnxruntime as ort
import numpy as np

# Load the ONNX model
session = ort.InferenceSession('net.onnx')

# Check input/output names and shapes
for inp in session.get_inputs():
    print(f"Input: {inp.name}, shape: {inp.shape}")
for out in session.get_outputs():
    print(f"Output: {out.name}, shape: {out.shape}")

# Run inference — ONNX Runtime uses NumPy arrays
x       = np.random.rand(4, 10).astype(np.float32)
outputs = session.run(['output'], {'input': x})
print(outputs[0].shape)   # (4, 3)

Dynamic vs Static Batch Size

Static batch (no dynamic_axes):
  Export with batch=1 → can only ever run with batch=1
  Fastest for fixed-size inference (e.g. embedded device)

Dynamic batch (dynamic_axes with 'batch_size'):
  Export once → run with any batch size
  Flexible for server deployment
  Slightly slower due to dynamic shape handling

Exporting a CNN

import torchvision.models as models
from torchvision.models import ResNet18_Weights

model = models.resnet18(weights=ResNet18_Weights.IMAGENET1K_V1)
model.eval()

dummy = torch.rand(1, 3, 224, 224)   # one RGB 224×224 image

torch.onnx.export(
    model,
    dummy,
    'resnet18.onnx',
    input_names=['image'],
    output_names=['class_scores'],
    dynamic_axes={'image': {0: 'batch'}, 'class_scores': {0: 'batch'}},
    opset_version=17
)
print("ResNet18 exported")

Comparing PyTorch vs ONNX Outputs

After exporting, verify that the ONNX model produces the same outputs as the original PyTorch model. Small floating-point differences are normal; large differences signal an export problem.

import numpy as np
import onnxruntime as ort

model.eval()
x = torch.rand(1, 10)

# PyTorch output
with torch.no_grad():
    pt_out = model(x).numpy()

# ONNX Runtime output
sess   = ort.InferenceSession('net.onnx')
ort_out = sess.run(['output'], {'input': x.numpy()})[0]

max_diff = np.abs(pt_out - ort_out).max()
print(f"Max absolute difference: {max_diff:.6f}")
# Should be < 1e-5 for correct export

Common Export Issues

Issue: "Unsupported op" during export
  Fix: Upgrade opset_version or rewrite the operation using supported ops

Issue: Outputs differ significantly from PyTorch
  Fix: Ensure model.eval() was called before export
       Check that dummy_input has the correct dtype (float32, not float64)

Issue: Dynamic shapes fail at runtime
  Fix: Declare all variable dimensions in dynamic_axes

Issue: Export fails on custom ops
  Fix: Register the custom op with ONNX or simplify the model

ONNX on Different Hardware

Target                   ONNX Runtime Provider
─────────────────────────────────────────────────────
CPU (any)                CPUExecutionProvider
NVIDIA GPU               CUDAExecutionProvider
Intel CPU/VPU            OpenVINOExecutionProvider
NVIDIA TensorRT          TensorrtExecutionProvider
DirectML (Windows GPU)   DmlExecutionProvider

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

Summary

ONNX export converts a PyTorch model to a framework-independent graph file that runs on ONNX Runtime, TensorRT, OpenVINO, and other engines. Call model.eval() before exporting. Provide a dummy input tensor matching the expected input shape. Use dynamic_axes to support variable batch sizes. Verify the export with onnx.checker.check_model and compare outputs between PyTorch and ONNX Runtime to confirm correctness. ONNX is the standard bridge between PyTorch research workflows and production inference systems on diverse hardware.