Chapter 3: ONNX - The Universal Format

WIA-AI-008 Standard • Estimated reading time: 50 minutes

3.1 What is ONNX?

The Open Neural Network Exchange (ONNX) is an open format designed to represent deep learning models in a framework-agnostic manner. Created in 2017 by Microsoft and Facebook (now Meta), ONNX has become the de facto standard for model interoperability, supported by all major frameworks including PyTorch, TensorFlow, Keras, MXNet, and many others.

The Philosophy Behind ONNX

ONNX embodies the 弘益人間 (Benefit All Humanity) philosophy by removing barriers between different AI ecosystems. It enables:

Researchers to share models regardless of the framework they used for training
Engineers to deploy models using optimized runtimes different from training frameworks
Organizations to avoid vendor lock-in and choose best-of-breed tools
The community to build a shared ecosystem of models and tools

3.2 ONNX Architecture

ONNX uses Protocol Buffers (protobuf) for serialization, defining models as directed acyclic graphs (DAGs).

Core Components

Graph: The computational graph containing nodes and edges
Nodes: Operations/operators (Conv, MatMul, ReLU, etc.)
Edges: Tensors flowing between nodes
Initializers: Constant tensors (weights, biases)
Metadata: Model version, producer info, documentation

# ONNX model structure
ModelProto
├── ir_version: int64
├── opset_import: [OpsetIdProto]
├── producer_name: string
├── producer_version: string
├── domain: string
├── model_version: int64
├── doc_string: string
├── graph: GraphProto
│   ├── node: [NodeProto]
│   ├── name: string
│   ├── input: [ValueInfoProto]
│   ├── output: [ValueInfoProto]
│   ├── initializer: [TensorProto]
│   └── value_info: [ValueInfoProto]
└── metadata_props: [StringStringEntryProto]

3.3 Exporting to ONNX

PyTorch to ONNX

PyTorch has excellent built-in ONNX export capabilities:

import torch
import torch.nn as nn

# Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
        self.relu = nn.ReLU()
        self.pool = nn.MaxPool2d(2)
        self.fc = nn.Linear(64 * 112 * 112, 10)

    def forward(self, x):
        x = self.pool(self.relu(self.conv1(x)))
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

model = SimpleModel()
model.eval()

# Prepare dummy input
dummy_input = torch.randn(1, 3, 224, 224)

# Export to ONNX
torch.onnx.export(
    model,                                # Model being exported
    dummy_input,                          # Model input (or tuple for multiple inputs)
    "simple_model.onnx",                 # Where to save the model
    export_params=True,                   # Store trained parameters
    opset_version=15,                    # ONNX version
    do_constant_folding=True,            # Optimize by constant folding
    input_names=['input'],               # Input names
    output_names=['output'],             # Output names
    dynamic_axes={                        # Dynamic dimensions
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

TensorFlow to ONNX

TensorFlow requires the tf2onnx tool:

# Install tf2onnx
# pip install tf2onnx

import tensorflow as tf
import tf2onnx

# Create a simple TensorFlow model
model = tf.keras.Sequential([
    tf.keras.layers.Conv2D(64, 3, padding='same', activation='relu', input_shape=(224, 224, 3)),
    tf.keras.layers.MaxPooling2D(2),
    tf.keras.layers.Flatten(),
    tf.keras.layers.Dense(10, activation='softmax')
])

# Save as SavedModel first
model.save('tf_model')

# Convert to ONNX via command line
# python -m tf2onnx.convert --saved-model tf_model --output model.onnx --opset 15

# Or programmatically
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
model_proto, _ = tf2onnx.convert.from_keras(model, input_signature=spec, opset=15)

with open("tf_model.onnx", "wb") as f:
    f.write(model_proto.SerializeToString())

Common Export Issues and Solutions

Issue	Cause	Solution
Unsupported operator	Custom op not in ONNX opset	Use supported alternatives or register custom op
Dynamic control flow	If/while loops in model	Use ONNX opset 13+ or refactor model
Shape inference failure	Ambiguous tensor shapes	Provide explicit shapes via dynamic_axes
Incorrect output values	Numerical precision issues	Check dtype conversions, use validation script

3.4 ONNX Runtime

ONNX Runtime is a high-performance inference engine optimized for ONNX models. It often outperforms native framework inference due to aggressive optimizations.

Basic Inference

import onnxruntime as ort
import numpy as np

# Create inference session
session = ort.InferenceSession("model.onnx")

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

# Prepare input data
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

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

print(f"Output shape: {outputs[0].shape}")
print(f"Predictions: {outputs[0]}")

Performance Optimization

import onnxruntime as ort

# Configure session options for maximum performance
session_options = ort.SessionOptions()

# Enable graph optimizations
session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL

# Set number of threads
session_options.intra_op_num_threads = 4
session_options.inter_op_num_threads = 4

# Enable profiling (for debugging)
session_options.enable_profiling = True

# Use specific execution providers (GPU, TensorRT, etc.)
providers = [
    'CUDAExecutionProvider',      # NVIDIA GPU
    'TensorrtExecutionProvider',  # TensorRT
    'CPUExecutionProvider'        # CPU fallback
]

session = ort.InferenceSession(
    "model.onnx",
    sess_options=session_options,
    providers=providers
)

# Check which provider is being used
print(f"Providers: {session.get_providers()}")

3.5 ONNX Model Inspection and Validation

Inspecting Model Structure

import onnx

# Load ONNX model
model = onnx.load("model.onnx")

# Check model validity
onnx.checker.check_model(model)
print("✓ Model is valid")

# Print model information
print(f"IR Version: {model.ir_version}")
print(f"Producer: {model.producer_name} {model.producer_version}")
print(f"Domain: {model.domain}")
print(f"Model Version: {model.model_version}")
print(f"Doc: {model.doc_string}")

# Inspect graph
graph = model.graph
print(f"\nGraph name: {graph.name}")

# Print inputs
print("\nInputs:")
for input in graph.input:
    print(f"  - {input.name}: {input.type.tensor_type.elem_type}, shape: {[d.dim_value for d in input.type.tensor_type.shape.dim]}")

# Print outputs
print("\nOutputs:")
for output in graph.output:
    print(f"  - {output.name}: {output.type.tensor_type.elem_type}")

# Print operators
print(f"\nNodes: {len(graph.node)}")
op_types = {}
for node in graph.node:
    op_types[node.op_type] = op_types.get(node.op_type, 0) + 1

print("\nOperator distribution:")
for op_type, count in sorted(op_types.items(), key=lambda x: -x[1]):
    print(f"  {op_type}: {count}")

Numerical Validation

import torch
import onnxruntime as ort
import numpy as np

# Original PyTorch model
model_pt = SimpleModel()
model_pt.eval()

# Export to ONNX
torch.onnx.export(model_pt, dummy_input, "model.onnx")

# Create ONNX Runtime session
session = ort.InferenceSession("model.onnx")

# Prepare test inputs
test_input = torch.randn(10, 3, 224, 224)

# PyTorch inference
with torch.no_grad():
    pt_output = model_pt(test_input).numpy()

# ONNX Runtime inference
ort_output = session.run(None, {
    session.get_inputs()[0].name: test_input.numpy()
})[0]

# Compare outputs
max_diff = np.max(np.abs(pt_output - ort_output))
mean_diff = np.mean(np.abs(pt_output - ort_output))

print(f"Max difference: {max_diff}")
print(f"Mean difference: {mean_diff}")

# Assert numerical equivalence (within tolerance)
assert np.allclose(pt_output, ort_output, rtol=1e-3, atol=1e-5), "Outputs don't match!"
print("✓ Numerical validation passed")

3.6 ONNX Operators and Opsets

ONNX defines a standardized set of operators. Different opset versions support different operators and features.

Opset Evolution

Opset	Released	Key Features
9	2019	Stable baseline, broad support
11	2020	Better RNN support, dynamic shapes
13	2021	Control flow (If, Loop), advanced ops
15	2022	Improved quantization, training support
17	2023	Better transformer support, LayerNorm

Common ONNX Operators

# Core operators
Conv, ConvTranspose                    # Convolution layers
Gemm, MatMul                           # Dense layers
Relu, Sigmoid, Tanh, Softmax          # Activations
BatchNormalization, LayerNormalization # Normalization
MaxPool, AveragePool, GlobalAveragePool # Pooling
Add, Sub, Mul, Div                     # Element-wise ops
Concat, Split, Reshape, Transpose      # Tensor manipulation
Gather, Scatter                        # Indexing operations
If, Loop                               # Control flow (opset 13+)
Cast, Clip, Pad                        # Utility operations

3.7 Custom Operators in ONNX

When ONNX doesn't support a specific operation, you can define custom operators:

import torch
from torch.onnx import register_custom_op_symbolic

# Define custom operator behavior
@torch.onnx.symbolic_helper.parse_args('v', 'v', 'f')
def custom_op(g, input1, input2, alpha):
    return g.op("custom_domain::CustomOp", input1, input2, alpha_f=alpha)

# Register the symbolic function
register_custom_op_symbolic('custom_domain::CustomOp', custom_op, 14)

# Use in model export
class ModelWithCustomOp(nn.Module):
    def forward(self, x, y):
        # Your custom operation
        return custom_operation(x, y, alpha=0.5)

# Export will now include the custom op
torch.onnx.export(model, (x, y), "model_custom.onnx")

3.8 ONNX Model Optimization

ONNX provides tools to optimize models for better performance:

from onnxruntime.transformers import optimizer
from onnxruntime.transformers.fusion_options import FusionOptions

# Load model
model_path = "model.onnx"

# Configure optimization options
opt_options = FusionOptions('bert')
opt_options.enable_gelu = True
opt_options.enable_layer_norm = True
opt_options.enable_attention = True
opt_options.enable_skip_layer_norm = True
opt_options.enable_embed_layer_norm = True
opt_options.enable_bias_skip_layer_norm = True
opt_options.enable_bias_gelu = True

# Optimize
optimized_model = optimizer.optimize_model(
    model_path,
    model_type='bert',
    num_heads=12,
    hidden_size=768,
    optimization_options=opt_options
)

# Save optimized model
optimized_model.save_model_to_file("model_optimized.onnx")

print(f"Original nodes: {len(onnx.load(model_path).graph.node)}")
print(f"Optimized nodes: {len(optimized_model.model.graph.node)}")

Graph Optimizations

Constant Folding: Evaluate constants at export time
Operator Fusion: Combine multiple ops (Conv + BatchNorm + ReLU)
Dead Code Elimination: Remove unused nodes
Shape Inference: Propagate tensor shapes through graph
Layout Transformation: Optimize NCHW vs NHWC layouts

3.9 ONNX for Different Model Types

Vision Models

# ResNet, EfficientNet, Vision Transformers
import torchvision.models as models

resnet50 = models.resnet50(pretrained=True)
resnet50.eval()

dummy_input = torch.randn(1, 3, 224, 224)
torch.onnx.export(
    resnet50,
    dummy_input,
    "resnet50.onnx",
    opset_version=15,
    input_names=['image'],
    output_names=['probabilities'],
    dynamic_axes={'image': {0: 'batch'}, 'probabilities': {0: 'batch'}}
)

NLP Models

# BERT, GPT, T5
from transformers import BertModel, BertTokenizer

tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertModel.from_pretrained('bert-base-uncased')
model.eval()

# Create dummy inputs
text = "Hello, ONNX!"
inputs = tokenizer(text, return_tensors='pt')

# Export with multiple inputs
torch.onnx.export(
    model,
    (inputs['input_ids'], inputs['attention_mask']),
    "bert.onnx",
    opset_version=14,
    input_names=['input_ids', 'attention_mask'],
    output_names=['last_hidden_state'],
    dynamic_axes={
        'input_ids': {0: 'batch', 1: 'sequence'},
        'attention_mask': {0: 'batch', 1: 'sequence'},
        'last_hidden_state': {0: 'batch', 1: 'sequence'}
    }
)

Generative Models

# GANs, VAEs, Diffusion Models
class SimpleGAN(nn.Module):
    def __init__(self):
        super().__init__()
        self.generator = nn.Sequential(
            nn.Linear(100, 256),
            nn.ReLU(),
            nn.Linear(256, 784),
            nn.Tanh()
        )

    def forward(self, z):
        return self.generator(z)

gan = SimpleGAN()
gan.eval()

latent = torch.randn(1, 100)
torch.onnx.export(
    gan,
    latent,
    "gan_generator.onnx",
    opset_version=15,
    input_names=['latent_vector'],
    output_names=['generated_image']
)

3.10 ONNX Ecosystem and Tools

Key Tools

onnx: Core library for ONNX model manipulation
onnxruntime: High-performance inference engine
onnx-simplifier: Simplifies and optimizes ONNX graphs
netron: Visual model viewer for ONNX and other formats
onnx-graphsurgeon: NVIDIA tool for graph manipulation

# Using onnx-simplifier
from onnxsim import simplify
import onnx

model = onnx.load("model.onnx")
simplified_model, check = simplify(model)

assert check, "Simplified model validation failed"
onnx.save(simplified_model, "model_simplified.onnx")

# Viewing with Netron
# Install: pip install netron
# Usage: netron model.onnx
# Opens interactive visualization in browser

Summary

ONNX is an open, framework-agnostic format for deep learning models
Uses Protocol Buffers to represent models as directed acyclic graphs
Core components: Graph, Nodes (operators), Edges (tensors), Initializers (weights)
Export from PyTorch via torch.onnx.export(), from TensorFlow via tf2onnx
ONNX Runtime provides optimized inference, often faster than native frameworks
Opset versions define supported operators; newer opsets add features
Custom operators can be registered for unsupported operations
Graph optimizations: constant folding, operator fusion, dead code elimination
Supports all model types: vision, NLP, generative, multimodal
Rich ecosystem: onnxruntime, onnx-simplifier, netron, graphsurgeon

Review Questions

What problem does ONNX solve in the machine learning ecosystem?
Describe the structure of an ONNX model. What are the main components?
How do you export a PyTorch model to ONNX? Include dynamic axes in your answer.
What is the difference between opset 9 and opset 15 in ONNX?
Why might ONNX Runtime be faster than native PyTorch or TensorFlow inference?
How do you validate that an ONNX model produces the same outputs as the original model?
What are common issues when exporting models to ONNX, and how can they be resolved?
Explain operator fusion and why it improves performance.
How would you handle a custom operation that's not in the ONNX opset?
Compare the process of exporting vision models vs NLP models to ONNX.