When developing modern artificial intelligence applications, efficiently and compatibly running trained machine learning models in production environments is a critical challenge. Different deep learning frameworks (PyTorch, TensorFlow, etc.) have their own standards for model formats, which can complicate cross-platform deployment and performance optimization. This is where ONNX (Open Neural Network Exchange) and ONNX Runtime come into play.

What is ONNX and Why is it Important?

ONNX is an open-source format standard for artificial intelligence models. This standard allows developers to easily transfer models from one framework to another and run them in different hardware/software environments. By converting a model to ONNX format, we make it independent of the framework it was trained in. This offers significant advantages, especially in the following scenarios:

• Consolidating models developed in different ML frameworks into a single format.
• Porting models to various deployment targets such as web, mobile, edge devices, and the cloud.
• Leveraging specialized accelerators (GPU, TPU, NPU, etc.) for performance optimization.

ONNX Runtime: High-Performance Inference Engine

ONNX Runtime is a cross-platform, high-performance inference engine that runs ONNX-formatted models. It can integrate with popular frameworks like PyTorch and TensorFlow and offers a wide range of hardware acceleration, from CPU to GPU (CUDA, TensorRT), and even specialized hardware (Intel OpenVINO, AMD ROCm). The main advantages of ONNX Runtime include:

• High Performance: Features built-in graph optimizations that shorten inference time and optimize resource utilization.
• Cross-Platform Support: Runs on many operating systems such as Windows, Linux, macOS, Android, iOS.
• Flexibility: Supports various hardware accelerators (Execution Providers).

Converting a PyTorch Model to ONNX Format with Python

Converting a PyTorch model to ONNX is quite straightforward. First, let's create and train a small PyTorch model (or use a pre-trained one).

import torch
import torch.nn as nn
import torch.optim as optim

# 1. Define a simple PyTorch model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.fc = nn.Linear(10, 2) # 10 inputs, 2 outputs
    
    def forward(self, x):
        return self.fc(x)

# Create the model and load weights (or train)
model = SimpleModel()
# Assume you trained your model or loaded pre-trained weights here
# For example, let's create random weights
dummy_input = torch.randn(1, 10) # Single sample, 10 features

# Set the model to evaluation mode
model.eval()

# 2. Convert the PyTorch model to ONNX
onnx_path = "simple_pytorch_model.onnx"
torch.onnx.export(
    model,
    dummy_input,
    onnx_path,
    export_params=True,
    opset_version=11, # Choose a supported opset version
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={'input': {0: 'batch_size'}, 'output': {0: 'batch_size'}} # Enable dynamic batch size
)

print(f"PyTorch model successfully converted to ONNX format: {onnx_path}")

Converting a TensorFlow/Keras Model to ONNX Format with Python

To convert TensorFlow or Keras models to ONNX, we can use the tf2onnx library. You may need to install this library first: pip install tf2onnx

import tensorflow as tf
import tf2onnx
from onnx.mapping import TENSOR_TYPE_TO_NP_TYPE

# 1. Define a simple Keras model
model = tf.keras.Sequential([
    tf.keras.layers.Dense(10, activation='relu', input_shape=(5,)), # 5 inputs
    tf.keras.layers.Dense(2, activation='softmax') # 2 outputs
])

# Train the model or load weights
# For example, let's compile the model
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# Check model output with random data (optional)
dummy_input = tf.random.normal((1, 5))
_ = model(dummy_input) # Ensure the model builds its structure

# 2. Convert the Keras model to ONNX
onnx_path = "simple_keras_model.onnx"
input_signature = [tf.TensorSpec(model.inputs[0].shape, model.inputs[0].dtype, name=model.inputs[0].name)]

model_proto, _ = tf2onnx.convert.from_keras(model, input_signature, opset=11, output_path=onnx_path)

print(f"Keras model successfully converted to ONNX format: {onnx_path}")

Performing Inference with ONNX Runtime

Let's see how to run the ONNX model we converted using ONNX Runtime. This applies to models converted from both PyTorch and TensorFlow.

import onnxruntime as ort
import numpy as np

# Load the ONNX model we saved earlier
onnx_model_path = "simple_pytorch_model.onnx" # or "simple_keras_model.onnx"

# Create an ONNX Runtime session
# If you want to use GPU: providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
# But you need to install ONNX Runtime for CUDA first (pip install onnxruntime-gpu)
sess_options = ort.SessionOptions()
session = ort.InferenceSession(onnx_model_path, sess_options, providers=['CPUExecutionProvider'])

# Get the input and output names of the model
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# Create random input data (in the expected shape and type of the model)
# For PyTorch model example (1, 10) float32
input_data = np.random.rand(1, 10).astype(np.float32) 
# For Keras model example (1, 5) float32
# input_data = np.random.rand(1, 5).astype(np.float32) 

# Perform inference
outputs = session.run([output_name], {input_name: input_data})

print("Inference results with ONNX Runtime:")
print(outputs[0])

Performance Tips and Best Practices

• Correct Opset Version: Ensure that the opset version you use during conversion is supported by your target ONNX Runtime version.
• Dynamic Input Shapes: If your model needs to support different batch sizes or variable input dimensions, specify this during ONNX conversion using the dynamic_axes parameter.
• Hardware Acceleration (Execution Providers): ONNX Runtime supports various hardware accelerators (Execution Providers). For example, using CUDAExecutionProvider or TensorRTExecutionProvider for NVIDIA GPUs can provide significant performance improvements over CPU. Remember to install the relevant libraries (onnxruntime-gpu) and specify the providers list when creating the session.
• Graph Optimizations: ONNX Runtime automatically optimizes model graphs. However, you can further enhance performance by applying additional optimizations or techniques like quantization for specific scenarios.
• Model Validation: Validate your converted ONNX model by checking if it produces the same output as the model in the original framework.

Conclusion

ONNX and ONNX Runtime are a powerful duo that simplify the deployment of machine learning models and enhance their performance. By using these tools, developers can break free from model dependencies, perform high-performance inference on different hardware platforms, and move AI applications to production faster and more efficiently. Learning these technologies is becoming an indispensable skill for software developers and AI enthusiasts when building modern AI infrastructures.