Chapter 2: ONNX Format Deep Dive

Open Neural Network Exchange - The Universal AI Format

弘益人間 · Benefit All Humanity

2.1 Introduction to ONNX

The Open Neural Network Exchange (ONNX) format was created in 2017 by Microsoft and Facebook (now Meta) as an open standard for representing machine learning models. ONNX has become the de facto interchange format for neural networks, enabling seamless model transfer between frameworks like PyTorch, TensorFlow, scikit-learn, and more.

ONNX is built on Protocol Buffers (protobuf), Google's language-neutral data serialization format. This choice provides efficient binary encoding, strong typing, and schema evolution capabilities essential for a production-grade format.

2.2 ONNX Model Structure

An ONNX model consists of several key components organized in a hierarchical structure:

Model Proto

The top-level container for the entire model:

message ModelProto {
  int64 ir_version = 1;           // ONNX IR version
  repeated OpsetImport opset_import = 8;  // Operator sets
  string producer_name = 2;        // Framework that created model
  string producer_version = 3;     // Framework version
  string domain = 4;               // Model domain
  int64 model_version = 5;         // Model version
  string doc_string = 6;           // Documentation
  GraphProto graph = 7;            // The computational graph
  repeated StringStringEntryProto metadata_props = 14;
}

Graph Proto

The computational graph defining the model's architecture:

message GraphProto {
  repeated NodeProto node = 1;           // Computation nodes
  string name = 2;                       // Graph name
  repeated TensorProto initializer = 5;  // Constant tensors (weights)
  string doc_string = 10;                // Documentation
  repeated ValueInfoProto input = 11;    // Input specifications
  repeated ValueInfoProto output = 12;   // Output specifications
  repeated ValueInfoProto value_info = 13; // Intermediate values
}

Node Proto

Individual operations in the computational graph:

message NodeProto {
  repeated string input = 1;       // Input tensor names
  repeated string output = 2;      // Output tensor names
  string name = 3;                 // Node name
  string op_type = 4;              // Operation type (Conv, Gemm, etc.)
  string domain = 7;               // Operator domain
  repeated AttributeProto attribute = 5;  // Operation attributes
  string doc_string = 6;           // Documentation
}

2.3 ONNX Operator Sets

ONNX defines operators through versioned operator sets (opsets). Each opset version introduces new operators or updates existing ones:

Opset	Release	Key Features
1-6	2017-2018	Basic operations, CNNs
7-10	2019	RNNs, control flow
11-13	2020-2021	Transformers, dynamic shapes
14-16	2022-2023	Quantization, optimization
17-18	2023-2024	LLM operations, efficiency

Common ONNX Operators

// Convolution
Conv: Convolution operation
  - dilations, group, kernel_shape, pads, strides
  - Input: X[N,C_in,H,W], W[C_out,C_in/group,kH,kW]
  - Output: Y[N,C_out,H_out,W_out]

// Matrix Multiplication
Gemm: General Matrix Multiplication (Y = α·AB + β·C)
  - alpha, beta, transA, transB

// Activation Functions
Relu, Sigmoid, Tanh, LeakyRelu, Elu, Selu, Swish

// Normalization
BatchNormalization, LayerNormalization, InstanceNormalization

// Pooling
MaxPool, AveragePool, GlobalMaxPool, GlobalAveragePool

// Attention (Opset 14+)
Attention, MultiHeadAttention

2.4 Tensor Representation in ONNX

ONNX tensors use the TensorProto message for constant values and weights:

message TensorProto {
  repeated int64 dims = 1;         // Tensor dimensions
  int32 data_type = 2;             // Element data type
  Segment segment = 3;             // For large tensors

  // Data storage (one of these):
  repeated float float_data = 4;   // FP32 values
  repeated int32 int32_data = 5;   // INT32 values
  bytes raw_data = 9;              // Raw binary data

  // External data reference
  repeated StringStringEntryProto external_data = 13;
  DataLocation data_location = 14;

  string name = 8;                 // Tensor name
  string doc_string = 12;          // Documentation
}

Data Types

Type	Value	Description
FLOAT	1	32-bit IEEE 754 floating point
UINT8	2	8-bit unsigned integer
INT8	3	8-bit signed integer
UINT16	4	16-bit unsigned integer
INT16	5	16-bit signed integer
INT32	6	32-bit signed integer
INT64	7	64-bit signed integer
FLOAT16	10	16-bit IEEE 754 half precision
BFLOAT16	16	16-bit brain float

2.5 Converting Models to ONNX

Most major frameworks provide tools to export models to ONNX format:

PyTorch to ONNX

import torch
import torch.onnx

# Load your trained model
model = MyModel()
model.load_state_dict(torch.load('model.pth'))
model.eval()

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

# Export to ONNX
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    export_params=True,
    opset_version=17,
    do_constant_folding=True,
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    }
)

TensorFlow to ONNX

import tensorflow as tf
import tf2onnx

# Load SavedModel
model = tf.saved_model.load('saved_model')

# Convert to ONNX
spec = (tf.TensorSpec((None, 224, 224, 3), tf.float32, name="input"),)
output_path = "model.onnx"

model_proto, _ = tf2onnx.convert.from_keras(
    model,
    input_signature=spec,
    opset=17,
    output_path=output_path
)

print(f"ONNX model saved to {output_path}")

Scikit-learn to ONNX

from skl2onnx import convert_sklearn
from skl2onnx.common.data_types import FloatTensorType

# Train scikit-learn model
from sklearn.ensemble import RandomForestClassifier
model = RandomForestClassifier(n_estimators=100)
model.fit(X_train, y_train)

# Define input type
initial_type = [('float_input', FloatTensorType([None, X_train.shape[1]]))]

# Convert to ONNX
onnx_model = convert_sklearn(
    model,
    initial_types=initial_type,
    target_opset=17
)

# Save
with open("sklearn_model.onnx", "wb") as f:
    f.write(onnx_model.SerializeToString())

2.6 ONNX Runtime

ONNX Runtime is a high-performance inference engine optimized for ONNX models:

Basic Inference

import onnxruntime as ort
import numpy as np

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

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

# Prepare input
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

# Configure session options
sess_options = ort.SessionOptions()

# Enable graph optimizations
sess_options.graph_optimization_level = \
    ort.GraphOptimizationLevel.ORT_ENABLE_ALL

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

# Enable profiling
sess_options.enable_profiling = True

# Create optimized session
session = ort.InferenceSession(
    "model.onnx",
    sess_options,
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']
)

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

2.7 Model Optimization in ONNX

ONNX provides various optimization techniques to improve inference performance:

Graph Optimization Levels

Quantization

from onnxruntime.quantization import quantize_dynamic, quantize_static

# Dynamic Quantization (no calibration data needed)
quantize_dynamic(
    model_input="model_fp32.onnx",
    model_output="model_int8.onnx",
    weight_type=QuantType.QInt8
)

# Static Quantization (requires calibration)
from onnxruntime.quantization import CalibrationDataReader

class DataReader(CalibrationDataReader):
    def __init__(self, calibration_data):
        self.data = calibration_data
        self.index = 0

    def get_next(self):
        if self.index < len(self.data):
            result = {"input": self.data[self.index]}
            self.index += 1
            return result
        return None

quantize_static(
    model_input="model_fp32.onnx",
    model_output="model_int8_static.onnx",
    calibration_data_reader=DataReader(calibration_dataset)
)

2.8 Custom Operators in ONNX

When standard ONNX operators don't cover your needs, you can define custom operators:

Defining a Custom Operator

from onnx import helper, TensorProto

# Define custom operator
custom_op = helper.make_node(
    'CustomOp',
    inputs=['X'],
    outputs=['Y'],
    domain='ai.onnx.custom',
    custom_attribute=42
)

# Register implementation in ONNX Runtime
import onnxruntime as ort

class CustomOpKernel:
    def __init__(self, custom_attribute):
        self.custom_attribute = custom_attribute

    def compute(self, x):
        # Your custom operation logic
        return x * self.custom_attribute

# Register kernel
ort.register_custom_ops_library("custom_ops.so")

2.9 ONNX Model Analysis Tools

Several tools help analyze and debug ONNX models:

Netron - Visual Model Explorer

Netron is a popular tool for visualizing ONNX models:

# Install Netron
pip install netron

# View model in browser
netron model.onnx

# Or use Python API
import netron
netron.start('model.onnx', port=8080)

ONNX Checker and Shape Inference

import onnx

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

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

# Infer shapes
model = onnx.shape_inference.infer_shapes(model)

# Print model structure
for node in model.graph.node:
    print(f"Node: {node.name}")
    print(f"  Op: {node.op_type}")
    print(f"  Inputs: {node.input}")
    print(f"  Outputs: {node.output}")

2.10 Best Practices for ONNX

Follow these guidelines when working with ONNX models:

Always Validate After Export

import torch
import onnxruntime as ort
import numpy as np

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

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

# Test input
test_input = np.random.randn(1, 3, 224, 224).astype(np.float32)

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

# ONNX output
onnx_output = session.run(None, {"input": test_input})[0]

# Compare
diff = np.abs(pt_output - onnx_output)
print(f"Max difference: {diff.max():.2e}")
print(f"Mean difference: {diff.mean():.2e}")

assert diff.max() < 1e-5, "Outputs differ significantly!"

Use Dynamic Axes for Flexibility

# Export with dynamic batch size and sequence length
torch.onnx.export(
    model,
    dummy_input,
    "model.onnx",
    dynamic_axes={
        'input': {0: 'batch', 1: 'sequence'},
        'output': {0: 'batch'}
    }
)

# Now model can handle any batch size and sequence length
session = ort.InferenceSession("model.onnx")
input_shape = session.get_inputs()[0].shape
print(input_shape)  # ['batch', 'sequence', 768]

Optimize for Target Hardware

# CPU optimization
session_cpu = ort.InferenceSession(
    "model.onnx",
    providers=['CPUExecutionProvider']
)

# GPU optimization (CUDA)
session_gpu = ort.InferenceSession(
    "model.onnx",
    providers=['CUDAExecutionProvider']
)

# TensorRT optimization (NVIDIA)
session_trt = ort.InferenceSession(
    "model.onnx",
    providers=['TensorrtExecutionProvider']
)

# Mobile optimization (CoreML for iOS)
from onnx_coreml import convert
coreml_model = convert(
    model="model.onnx",
    minimum_ios_deployment_target='13'
)