How to export PyTorch models to ONNX and run inference with ONNX Runtime?

Deploying PyTorch models directly into production environments often introduces heavy Python runtime dependencies and hardware abstraction overhead. Converting PyTorch neural networks into Open Neural Network Exchange format allows optimized execution across heterogeneous platforms using ONNX Runtime.

Model Serialization and Execution

Below is a working workflow illustrating how to export a computer vision or classification PyTorch module into ONNX format and execute inference sessions.

import torch
import torch.nn as nn
import onnxruntime as ort
import numpy as np

# Define a simple linear model module
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(10, 2)
        
    def forward(self, x):
        return self.linear(x)

# Instantiate model and create dummy input tensor matching input shape
pytorch_model = SimpleModel()
pytorch_model.eval()
dummy_input = torch.randn(1, 10)

# Export PyTorch model graph to ONNX file
torch.onnx.export(
    pytorch_model,
    dummy_input,
    "model.onnx",
    input_names=["input"],
    output_names=["output"],
    dynamic_axes={"input": {0: "batch_size"}, "output": {0: "batch_size"}}
)

# Load exported ONNX model into ONNX Runtime session
ort_session = ort.InferenceSession("model.onnx")

# Run inference using numpy array format
ort_inputs = {"input": dummy_input.numpy()}
ort_outputs = ort_session.run(None, ort_inputs)
print("ONNX Inference output shape:", ort_outputs[0].shape)

Technical Questions:

• How do you resolve custom operator conversion errors when trying to export PyTorch models containing dynamic loops or custom C++ extensions?
• What graph optimization settings in ONNX Runtime yield the highest throughput improvements on CPU edge devices?