AI篇-ONNX学习指南：从原理到实战，打通模型部署的"任督二脉"

帅气的你
73 阅读23分钟

AI篇-ONNX学习指南：从原理到实战，打通模型部署的"任督二脉"

引言

在深度学习模型从训练到部署的完整链路中，开发者常常面临这样的困境：

  • 框架壁垒：PyTorch训练的模型无法直接在TensorFlow Serving上运行，TensorFlow的模型也无法直接用于PyTorch Mobile
  • 跨平台兼容性差：同一模型需要在不同平台（CPU、GPU、移动端、边缘设备）部署时，需要针对每个平台重新适配和优化
  • 推理性能不足：训练框架的推理引擎往往不是为生产环境优化的，导致推理速度慢、资源占用高

ONNX（Open Neural Network Exchange）应运而生，它是由微软、Facebook、亚马逊等公司共同推出的开放神经网络交换格式。ONNX的核心价值在于：

  1. 统一模型格式：将不同框架训练的模型转换为统一的ONNX格式，实现"一次转换，多端部署"
  2. 跨框架兼容：支持PyTorch、TensorFlow、Keras、MXNet、Caffe2等主流框架的模型转换
  3. 高性能推理：ONNX Runtime作为专门的推理引擎，针对不同硬件平台进行了深度优化
  4. 生态完善：拥有丰富的工具链和社区支持，覆盖模型转换、优化、可视化、部署等全流程

本文将从ONNX的核心概念入手，通过丰富的实战案例，帮助读者系统掌握ONNX的使用方法，最终能够独立完成模型的ONNX转换、优化与多平台部署。

学习目标：

  • 理解ONNX的技术架构与核心优势
  • 掌握主流框架模型转ONNX的完整流程
  • 学会ONNX模型的优化与量化方法
  • 实现ONNX模型在多平台的推理与部署
  • 能够解决ONNX使用过程中的常见问题

第一章：ONNX核心概念深度解析

1.1 什么是ONNX？

ONNX（Open Neural Network Exchange）是一个开放的、跨平台的深度学习模型表示格式标准。它的本质是一种中间表示（Intermediate Representation, IR），类似于编程语言中的"字节码"。

官方定义：

ONNX provides an open source format for AI models, both deep learning and traditional ML. It defines an extensible computation graph model, as well as definitions of built-in operators and standard data types.

核心特点：

  • 格式开放：基于Protobuf序列化，格式规范公开透明
  • 框架无关：不依赖任何特定的深度学习框架
  • 平台无关：可以在CPU、GPU、移动端、边缘设备等多种平台上运行

1.2 ONNX的核心优势

1.2.1 跨框架支持

ONNX支持主流深度学习框架的模型转换：

框架转换工具支持程度
PyTorchtorch.onnx原生支持，支持动态图
TensorFlowtf2onnx支持TF 1.x和2.x
Kerastf2onnx通过TensorFlow转换
MXNetmxnet-onnx官方支持
Caffe2内置支持原生支持
PaddlePaddlepaddle2onnx官方支持
MindSporemindspore-onnx社区支持
1.2.2 跨平台部署

ONNX模型可以在多种硬件平台上运行：

  • CPU：通过ONNX Runtime CPU执行提供者
  • GPU：支持CUDA、TensorRT、DirectML等
  • 移动端：ONNX Runtime Mobile支持Android/iOS
  • 边缘设备：支持树莓派、Jetson等嵌入式设备
1.2.3 生态完善

ONNX拥有丰富的工具链：

  • 模型转换：各框架官方/社区转换工具
  • 模型优化：onnxoptimizer、onnx-simplifier
  • 模型可视化：Netron、ONNX GraphSurgeon
  • 推理引擎：ONNX Runtime、TensorRT、OpenVINO
  • 性能分析：ONNX Runtime Profiler
1.2.4 支持动态图与静态图
  • 静态图：输入输出维度固定，推理性能最优
  • 动态图：支持动态输入维度（通过dynamic_axes参数），灵活性更高

1.3 ONNX的技术架构

1.3.1 计算图模型（Computation Graph）

ONNX使用**有向无环图（DAG）**来表示神经网络：

输入节点 → 算子节点 → 算子节点 → ... → 输出节点

图结构组成：

  • 节点（Node）：表示一个算子（如Conv、Relu、MatMul）
  • 边（Edge）：表示数据流（Tensor）
  • 输入/输出：图的入口和出口
1.3.2 算子集（ONNX Operator Set）

ONNX定义了标准的算子集合，每个算子都有明确的输入输出规范。ONNX Operator Set版本不断演进，新版本会添加更多算子支持。

常用算子类别：

  • 数学运算：Add、Mul、Sub、Div、Pow
  • 神经网络层：Conv、MaxPool、AveragePool、BatchNormalization
  • 激活函数：Relu、Sigmoid、Tanh、Gelu
  • 张量操作：Reshape、Transpose、Concat、Split
  • 循环网络：LSTM、GRU、RNN
1.3.3 Tensor数据类型

ONNX支持多种数据类型：

数据类型说明常用场景
FLOAT32位浮点数标准模型推理
FLOAT1616位浮点数模型量化，减少内存
INT88位整数模型量化，极致压缩
INT3232位整数索引、形状信息
INT6464位整数大范围索引
BOOL布尔类型条件判断
STRING字符串文本处理

1.4 ONNX与ONNX Runtime的关系

重要区分：

  • ONNX：是模型格式标准，定义了模型的存储格式和算子规范
  • ONNX Runtime：是推理引擎，负责加载ONNX模型并在不同硬件平台上执行推理

关系类比：

  • ONNX = PDF文件格式
  • ONNX Runtime = PDF阅读器

ONNX Runtime的作用：

  1. 模型加载：解析ONNX格式文件，构建计算图
  2. 图优化：执行图级别的优化（算子融合、常量折叠等）
  3. 执行调度：根据硬件平台选择最优的执行提供者（Provider）
  4. 内存管理：高效管理推理过程中的内存分配

ONNX Runtime执行提供者（Providers）：

  • CPUExecutionProvider：CPU推理（默认）
  • CUDAExecutionProvider：NVIDIA GPU推理
  • TensorrtExecutionProvider：TensorRT加速
  • OpenVINOExecutionProvider：Intel硬件加速
  • DirectMLExecutionProvider：Windows DirectML

第二章：环境搭建与基础工具准备

2.1 软硬件环境要求

2.1.1 硬件要求
  • CPU：支持x86_64或ARM架构
  • 内存：建议8GB以上
  • GPU（可选）：NVIDIA GPU（支持CUDA 11.0+），用于GPU加速推理
2.1.2 软件要求
  • 操作系统：Linux、macOS、Windows
  • Python版本：Python 3.7 - 3.11（推荐3.8或3.9）
  • 深度学习框架：PyTorch 1.8+ 或 TensorFlow 2.x

2.2 环境安装

2.2.1 安装ONNX核心库
# 安装ONNX（模型格式支持）
pip install onnx

# 安装ONNX Runtime（推理引擎）
# CPU版本
pip install onnxruntime

# GPU版本（需要CUDA支持）
pip install onnxruntime-gpu

# 验证安装
python -c "import onnx; import onnxruntime; print('ONNX version:', onnx.__version__); print('ONNX Runtime version:', onnxruntime.__version__)"
2.2.2 安装深度学习框架

PyTorch安装：

# CPU版本
pip install torch torchvision torchaudio

# GPU版本（CUDA 11.8）
pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118

TensorFlow安装：

# TensorFlow 2.x
pip install tensorflow

# 验证安装
python -c "import torch; import tensorflow as tf; print('PyTorch:', torch.__version__); print('TensorFlow:', tf.__version__)"
2.2.3 安装辅助工具
# ONNX模型优化工具
pip install onnxoptimizer

# ONNX模型简化工具
pip install onnx-simplifier

# ONNX GraphSurgeon（模型修改工具）
pip install onnx-graphsurgeon --extra-index-url https://pypi.ngc.nvidia.com

# 模型可视化工具Netron（可选，有Web版本和桌面版）
# Web版本：访问 https://netron.app
# 桌面版：从 https://github.com/lutzroeder/netron/releases 下载

2.3 常用工具介绍

2.3.1 Netron - 模型可视化工具

功能：可视化ONNX模型的计算图结构，查看节点、输入输出、权重等信息。

使用方法：

  1. Web版本：

    • 访问 netron.app
    • 点击"Open Model"上传ONNX模型文件
    • 即可查看模型结构

  2. 桌面版本：

    # 下载对应平台的安装包
# macOS
brew install netron

# 或直接运行
netron model.onnx

可视化内容：

  • 计算图结构（节点、边）
  • 每个节点的输入输出维度
  • 权重参数信息
  • 算子类型和属性
2.3.2 ONNX GraphSurgeon - 模型修改工具

功能：在Python代码中修改ONNX模型结构（添加/删除节点、修改输入输出等）。

基础使用示例：

import onnx
import onnx_graphsurgeon as gs

# 加载ONNX模型
model = onnx.load("model.onnx")
graph = gs.import_onnx(model)

# 查找节点
nodes = graph.nodes
for node in nodes:
    print(f"Node: {node.name}, Op: {node.op}, Outputs: {node.outputs}")

# 修改节点（示例：重命名输出）
graph.outputs[0].name = "new_output_name"

# 导出修改后的模型
onnx.save(gs.export_onnx(graph), "modified_model.onnx")
2.3.3 ONNX Checker - 模型验证工具

功能：检查ONNX模型是否符合规范。

import onnx

# 加载模型
model = onnx.load("model.onnx")

# 检查模型
onnx.checker.check_model(model)
print("Model is valid!")

第三章：核心实战——主流框架模型转ONNX格式

3.1 实战1：PyTorch模型转ONNX

步骤1：准备预训练PyTorch模型

我们以ResNet18为例，展示完整的转换流程：

import torch
import torchvision.models as models
import onnx

# 加载预训练模型
model = models.resnet18(pretrained=True)
model.eval()  # 设置为评估模式

# 或者使用自定义模型
class SimpleClassifier(torch.nn.Module):
    def __init__(self):
        super(SimpleClassifier, self).__init__()
        self.conv1 = torch.nn.Conv2d(3, 64, kernel_size=3, stride=1, padding=1)
        self.relu = torch.nn.ReLU()
        self.pool = torch.nn.AdaptiveAvgPool2d((1, 1))
        self.fc = torch.nn.Linear(64, 10)
    
    def forward(self, x):
        x = self.conv1(x)
        x = self.relu(x)
        x = self.pool(x)
        x = x.view(x.size(0), -1)
        x = self.fc(x)
        return x

# model = SimpleClassifier()
步骤2：使用torch.onnx.export()完成转换

基础转换：

# 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 导出ONNX模型
onnx_path = "resnet18.onnx"
torch.onnx.export(
    model,                          # 要转换的模型
    dummy_input,                    # 模型输入（用于追踪模型结构）
    onnx_path,                      # 输出ONNX文件路径
    input_names=['input'],          # 输入名称
    output_names=['output'],        # 输出名称
    verbose=False                   # 是否打印详细信息
)

print(f"Model exported to {onnx_path}")

关键参数详解：

参数说明示例
model要转换的PyTorch模型model
args模型输入（可以是tuple/list）dummy_input 或 (dummy_input,)
f输出文件路径或文件对象"model.onnx"
input_names输入节点名称列表['input']
output_names输出节点名称列表['output']
dynamic_axes动态维度设置（见下文）{'input': {0: 'batch_size'}}
opset_versionONNX算子集版本11（默认）
do_constant_folding是否进行常量折叠优化True（默认）
verbose是否打印详细信息False
步骤3：动态维度设置

问题：默认情况下，ONNX模型的输入输出维度是固定的。如果训练时输入是(1, 3, 224, 224)，转换后的模型只能接受这个固定尺寸的输入。

解决方案：使用dynamic_axes参数设置动态维度。

# 设置动态batch size和图像尺寸
torch.onnx.export(
    model,
    dummy_input,
    "resnet18_dynamic.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size', 2: 'height', 3: 'width'},  # 第0维是batch，第2、3维是图像高宽
        'output': {0: 'batch_size'}  # 输出也支持动态batch
    }
)

动态维度说明：

  • {0: 'batch_size'}：第0维（batch维度）是动态的，可以接受任意batch size
  • {2: 'height', 3: 'width'}：图像的高和宽可以是动态的
步骤4：转换后模型验证
import onnx

# 加载转换后的模型
onnx_model = onnx.load("resnet18.onnx")

# 检查模型有效性
try:
    onnx.checker.check_model(onnx_model)
    print("✓ Model is valid!")
except onnx.checker.ValidationError as e:
    print(f"✗ Model is invalid: {e}")

# 查看模型信息
print(f"Model inputs: {[input.name for input in onnx_model.graph.input]}")
print(f"Model outputs: {[output.name for output in onnx_model.graph.output]}")
print(f"ONNX opset version: {onnx_model.opset_import[0].version}")

完整转换示例：

import torch
import torchvision.models as models
import onnx

# 1. 加载模型
model = models.resnet18(pretrained=True)
model.eval()

# 2. 创建示例输入
dummy_input = torch.randn(1, 3, 224, 224)

# 3. 导出ONNX（支持动态batch）
torch.onnx.export(
    model,
    dummy_input,
    "resnet18.onnx",
    input_names=['input'],
    output_names=['output'],
    dynamic_axes={
        'input': {0: 'batch_size'},
        'output': {0: 'batch_size'}
    },
    opset_version=11,
    do_constant_folding=True
)

# 4. 验证模型
onnx_model = onnx.load("resnet18.onnx")
onnx.checker.check_model(onnx_model)
print("✓ Conversion successful!")

3.2 实战2：TensorFlow/Keras模型转ONNX

步骤1：准备Keras模型
import tensorflow as tf
from tensorflow import keras

# 加载预训练模型（以MobileNetV2为例）
model = keras.applications.MobileNetV2(
    input_shape=(224, 224, 3),
    weights='imagenet',
    include_top=True
)

# 或者构建自定义模型
# model = keras.Sequential([
#     keras.layers.Conv2D(32, 3, activation='relu', input_shape=(224, 224, 3)),
#     keras.layers.GlobalAveragePooling2D(),
#     keras.layers.Dense(1000, activation='softmax')
# ])
步骤2：使用tf2onnx完成转换

安装tf2onnx：

pip install tf2onnx

转换代码：

import tf2onnx
import onnx

# 方法1：使用tf2onnx.convert（推荐）
onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    output_path="mobilenetv2.onnx",
    opset=11,
    input_signature=None  # 如果模型有动态输入，需要指定input_signature
)

# 方法2：使用命令行工具
# 先保存SavedModel格式
model.save("saved_model/mobilenetv2")

# 然后使用命令行转换
# python -m tf2onnx.convert --saved-model saved_model/mobilenetv2 --output mobilenetv2.onnx --opset 11

完整转换示例：

import tensorflow as tf
from tensorflow import keras
import tf2onnx
import onnx

# 1. 加载或构建模型
model = keras.applications.MobileNetV2(
    input_shape=(224, 224, 3),
    weights='imagenet'
)

# 2. 转换为ONNX
onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    output_path="mobilenetv2.onnx",
    opset=11
)

# 3. 验证模型
onnx.checker.check_model(onnx_model)
print("✓ Conversion successful!")
步骤3：转换常见问题与解决方案

问题1：算子不支持

错误信息：

RuntimeError: Unsupported ONNX opset version: 12

解决方案：

# 降低opset版本
onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    output_path="model.onnx",
    opset=11  # 使用较低的opset版本
)

问题2：动态输入维度

解决方案：

# 指定input_signature
import tensorflow as tf

# 定义动态输入签名
input_signature = [tf.TensorSpec([None, 224, 224, 3], tf.float32, name='input')]

onnx_model, _ = tf2onnx.convert.from_keras(
    model,
    output_path="model.onnx",
    opset=11,
    input_signature=input_signature
)

问题3：版本不匹配

解决方案：

# 确保版本兼容
pip install tensorflow==2.10.0  # 使用稳定版本
pip install tf2onnx==1.13.0

3.3 补充：其他框架转ONNX

MindSpore转ONNX
# 安装转换工具
# pip install mindspore

# 转换代码
import mindspore as ms
from mindspore import export

# 加载MindSpore模型
model = ms.load_checkpoint("model.ckpt")

# 导出ONNX
export(model, ms.Tensor(np.random.randn(1, 3, 224, 224).astype(np.float32)), 
       file_name="model", file_format="ONNX")
PaddlePaddle转ONNX
# 安装转换工具
pip install paddle2onnx

# 使用命令行转换
paddle2onnx --model_dir paddle_model \
            --model_filename inference.pdmodel \
            --params_filename inference.pdiparams \
            --save_file model.onnx \
            --opset_version 11

第四章：ONNX模型优化与量化

4.1 为什么需要模型优化？

优化目标：

  1. 减小模型体积：降低存储和传输成本
  2. 提升推理速度：减少计算时间，提高吞吐量
  3. 降低内存占用：适配资源受限的设备

优化方法分类：

  • 图优化：消除冗余节点、算子融合、常量折叠
  • 量化：降低数值精度（FP32 → FP16/INT8）
  • 剪枝：移除不重要的权重或通道
  • 蒸馏：用大模型指导小模型训练

4.2 ONNX原生优化

4.2.1 使用onnxoptimizer进行图优化

安装：

pip install onnxoptimizer

优化代码：

import onnx
from onnxoptimizer import optimize_model

# 加载原始模型
model = onnx.load("model.onnx")

# 应用优化
# 可用的优化级别：'basic', 'extended', 'no_optimization'
optimized_model = optimize_model(
    model,
    passes=[  # 指定要应用的优化pass
        'eliminate_deadend',
        'eliminate_identity',
        'eliminate_nop_transpose',
        'eliminate_nop_pad',
        'fuse_matmul_add_bias_into_gemm',
        'fuse_bn_into_conv',
        'fuse_consecutive_concats',
        'fuse_consecutive_log_softmax',
        'fuse_consecutive_reduce_unsqueeze',
        'fuse_consecutive_squeezes',
        'fuse_transpose_into_gemm',
    ]
)

# 保存优化后的模型
onnx.save(optimized_model, "model_optimized.onnx")
print("✓ Model optimized!")

常用优化Pass说明：

Pass名称功能
eliminate_deadend消除死节点（无输出的节点）
eliminate_identity消除恒等变换节点
fuse_bn_into_conv将BatchNorm融合到Conv中
fuse_matmul_add_bias_into_gemm将MatMul+Add融合为GEMM
fuse_consecutive_concats融合连续的Concat操作
4.2.2 使用onnx-simplifier简化模型

安装：

pip install onnx-simplifier

简化代码：

import onnx
from onnxsim import simplify

# 加载模型
model = onnx.load("model.onnx")

# 简化模型
simplified_model, check = simplify(model)

# 检查简化是否成功
assert check, "Simplified ONNX model could not be validated"

# 保存简化后的模型
onnx.save(simplified_model, "model_simplified.onnx")
print("✓ Model simplified!")

命令行使用：

onnxsim model.onnx model_simplified.onnx

4.3 模型量化

4.3.1 量化的核心原理

量化定义：将模型中的浮点权重和激活值转换为低精度整数，从而减少模型大小和加速推理。

量化类型：

量化类型说明精度损失速度提升
FP32原始精度无基准
FP16半精度浮点极小1.5-2x
INT88位整数较小2-4x

量化方法：

  • 动态量化：权重量化，激活值在推理时动态量化
  • 静态量化：权重和激活值都预先量化，需要校准数据集
  • 量化感知训练（QAT）：在训练过程中模拟量化，精度损失最小
4.3.2 使用ONNX Runtime进行量化

动态量化（Dynamic Quantization）：

import onnx
from onnxruntime.quantization import quantize_dynamic, QuantType

# 加载FP32模型
model_fp32 = "model.onnx"

# 动态量化（权重INT8，激活值FP32）
model_int8 = "model_int8_dynamic.onnx"
quantize_dynamic(
    model_fp32,
    model_int8,
    weight_type=QuantType.QUINT8  # 权重量化类型
)

print("✓ Dynamic quantization completed!")

静态量化（Static Quantization）：

import onnx
from onnxruntime.quantization import quantize_static, CalibrationDataReader, QuantType

# 定义校准数据读取器
class DataReader(CalibrationDataReader):
    def __init__(self, calibration_dataset):
        self.dataset = calibration_dataset
        self.iter = iter(calibration_dataset)
    
    def get_next(self):
        try:
            return {"input": next(self.iter)}  # 返回字典，key为输入名称
        except StopIteration:
            return None

# 准备校准数据集（100-200个样本即可）
calibration_dataset = [np.random.randn(1, 3, 224, 224).astype(np.float32) 
                       for _ in range(100)]
data_reader = DataReader(calibration_dataset)

# 静态量化
model_fp32 = "model.onnx"
model_int8 = "model_int8_static.onnx"
quantize_static(
    model_fp32,
    model_int8,
    data_reader,
    quant_type=QuantType.QInt8,  # 使用INT8量化
    activation_type=QuantType.QUInt8,
    weight_type=QuantType.QInt8
)

print("✓ Static quantization completed!")
4.3.3 量化前后性能对比

性能对比脚本：

import onnxruntime as ort
import numpy as np
import time

# 加载模型
model_fp32 = "model.onnx"
model_int8 = "model_int8_static.onnx"

# 创建推理会话
session_fp32 = ort.InferenceSession(model_fp32)
session_int8 = ort.InferenceSession(model_int8, 
                                     providers=['CPUExecutionProvider'])

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)
input_name = session_fp32.get_inputs()[0].name

# 测试FP32模型
start = time.time()
for _ in range(100):
    _ = session_fp32.run(None, {input_name: input_data})
fp32_time = (time.time() - start) / 100

# 测试INT8模型
start = time.time()
for _ in range(100):
    _ = session_int8.run(None, {input_name: input_data})
int8_time = (time.time() - start) / 100

# 计算模型大小
import os
fp32_size = os.path.getsize(model_fp32) / (1024 * 1024)  # MB
int8_size = os.path.getsize(model_int8) / (1024 * 1024)  # MB

print(f"FP32 Model:")
print(f"  Size: {fp32_size:.2f} MB")
print(f"  Inference Time: {fp32_time*1000:.2f} ms")
print(f"\nINT8 Model:")
print(f"  Size: {int8_size:.2f} MB")
print(f"  Inference Time: {int8_time*1000:.2f} ms")
print(f"\nSpeedup: {fp32_time/int8_time:.2f}x")
print(f"Size Reduction: {(1-int8_size/fp32_size)*100:.1f}%")

典型结果：

  • 模型大小：减少约75%（FP32 → INT8）
  • 推理速度：提升2-4倍（取决于硬件和模型）
  • 精度损失：通常<1%（在ImageNet等数据集上）

第五章：ONNX模型推理与多平台部署

5.1 基础推理：使用ONNX Runtime进行模型推理

步骤1：加载ONNX模型
import onnxruntime as ort
import numpy as np

# 创建推理会话
session = ort.InferenceSession("model.onnx")

# 查看模型输入输出信息
for input_meta in session.get_inputs():
    print(f"Input: {input_meta.name}, Shape: {input_meta.shape}, Type: {input_meta.type}")

for output_meta in session.get_outputs():
    print(f"Output: {output_meta.name}, Shape: {output_meta.shape}, Type: {output_meta.type}")
步骤2：准备输入数据（预处理）

图像预处理示例：

from PIL import Image
import numpy as np

def preprocess_image(image_path, target_size=(224, 224)):
    """
    图像预处理：与训练时保持一致
    """
    # 1. 加载图像
    img = Image.open(image_path).convert('RGB')
    
    # 2. 调整大小
    img = img.resize(target_size)
    
    # 3. 转换为numpy数组并归一化
    img_array = np.array(img).astype(np.float32) / 255.0
    
    # 4. 标准化（ImageNet均值和标准差）
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img_array = (img_array - mean) / std
    
    # 5. 调整维度顺序：HWC -> CHW
    img_array = img_array.transpose(2, 0, 1)
    
    # 6. 添加batch维度：CHW -> BCHW
    img_array = np.expand_dims(img_array, axis=0)
    
    return img_array

# 使用示例
input_data = preprocess_image("test_image.jpg")
步骤3：执行推理并解析输出
# 获取输入输出名称
input_name = session.get_inputs()[0].name
output_name = session.get_outputs()[0].name

# 执行推理
outputs = session.run([output_name], {input_name: input_data})

# 解析输出（分类任务示例）
predictions = outputs[0]
predicted_class = np.argmax(predictions, axis=1)
confidence = np.max(predictions, axis=1)

print(f"Predicted class: {predicted_class[0]}")
print(f"Confidence: {confidence[0]:.4f}")

完整推理示例：

import onnxruntime as ort
import numpy as np
from PIL import Image

# 1. 加载模型
session = ort.InferenceSession("resnet18.onnx")

# 2. 预处理图像
def preprocess(image_path):
    img = Image.open(image_path).convert('RGB').resize((224, 224))
    img_array = np.array(img).astype(np.float32) / 255.0
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    img_array = (img_array - mean) / std
    img_array = img_array.transpose(2, 0, 1)
    return np.expand_dims(img_array, axis=0)

# 3. 执行推理
input_data = preprocess("test_image.jpg")
input_name = session.get_inputs()[0].name
outputs = session.run(None, {input_name: input_data})

# 4. 后处理
predictions = outputs[0]
predicted_class = np.argmax(predictions, axis=1)[0]
print(f"Predicted class: {predicted_class}")
步骤4：对比原框架与ONNX Runtime的推理速度
import torch
import onnxruntime as ort
import numpy as np
import time

# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)

# PyTorch推理
model_pytorch = torch.jit.load("model.pt")  # 或加载原始PyTorch模型
model_pytorch.eval()

start = time.time()
with torch.no_grad():
    output_pytorch = model_pytorch(torch.from_numpy(input_data))
pytorch_time = time.time() - start

# ONNX Runtime推理
session = ort.InferenceSession("model.onnx")
input_name = session.get_inputs()[0].name

start = time.time()
output_onnx = session.run(None, {input_name: input_data})
onnx_time = time.time() - start

print(f"PyTorch Inference Time: {pytorch_time*1000:.2f} ms")
print(f"ONNX Runtime Inference Time: {onnx_time*1000:.2f} ms")
print(f"Speedup: {pytorch_time/onnx_time:.2f}x")

5.2 多平台部署实战

5.2.1 部署到CPU/GPU

指定执行提供者（Provider）：

import onnxruntime as ort

# CPU推理（默认）
session_cpu = ort.InferenceSession(
    "model.onnx",
    providers=['CPUExecutionProvider']
)

# GPU推理（需要安装onnxruntime-gpu）
session_gpu = ort.InferenceSession(
    "model.onnx",
    providers=['CUDAExecutionProvider', 'CPUExecutionProvider']  # 优先使用CUDA，失败则回退到CPU
)

# 查看可用的提供者
print("Available providers:", ort.get_available_providers())

性能优化配置：

# 配置会话选项
options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL  # 启用图优化
options.intra_op_num_threads = 4  # 设置线程数
options.inter_op_num_threads = 4

session = ort.InferenceSession(
    "model.onnx",
    sess_options=options,
    providers=['CPUExecutionProvider']
)
5.2.2 部署到移动端

ONNX Runtime Mobile：

  1. Android部署：

    • 使用ONNX Runtime Mobile的Java/Kotlin API
    • 将ONNX模型打包到APK中
    • 在Android应用中加载和推理

  2. iOS部署：

    • 使用ONNX Runtime Mobile的Objective-C/Swift API
    • 将ONNX模型添加到Xcode项目
    • 在iOS应用中加载和推理

示例（Android）：

// Java代码示例
import ai.onnxruntime.*;

// 创建推理环境
OrtEnvironment env = OrtEnvironment.getEnvironment();
OrtSession.SessionOptions opts = new OrtSession.SessionOptions();

// 加载模型
OrtSession session = env.createSession("model.onnx", opts);

// 准备输入
OnnxTensor inputTensor = OnnxTensor.createTensor(env, inputData);
Map<String, OnnxTensor> inputs = Collections.singletonMap("input", inputTensor);

// 执行推理
OrtSession.Result result = session.run(inputs);
5.2.3 部署到边缘设备

与TensorRT结合（NVIDIA Jetson）：

import onnxruntime as ort

# 使用TensorRT执行提供者
session = ort.InferenceSession(
    "model.onnx",
    providers=['TensorrtExecutionProvider', 'CUDAExecutionProvider', 'CPUExecutionProvider']
)

与OpenVINO结合（Intel硬件）：

import onnxruntime as ort

# 使用OpenVINO执行提供者
session = ort.InferenceSession(
    "model.onnx",
    providers=['OpenVINOExecutionProvider', 'CPUExecutionProvider']
)

树莓派部署：

# 在树莓派上安装ONNX Runtime
# 注意：需要ARM架构的wheel包
pip install onnxruntime

# Python代码与CPU推理相同

第六章：常见问题与解决方案

6.1 转换失败类问题

问题1：算子不支持

错误信息：

RuntimeError: ONNX export failed: Couldn't export operator aten::xxx

解决方案：

  1. 检查算子支持情况：

    # 查看PyTorch版本支持的ONNX算子
import torch
print(torch.onnx.export.__doc__)

  2. 使用自定义算子映射：

    # 对于不支持的算子，可以注册自定义转换函数
import torch.onnx.symbolic_registry as registry

# 示例：自定义算子转换（需要根据具体算子实现）
# @registry.register("aten::custom_op", "", 11)
# def custom_op(g, input):
#     return g.op("CustomOp", input)

  3. 替换不支持的算子：

    # 在模型中用支持的算子替换不支持的算子
# 例如：用Conv+BN融合替换某些复杂操作
问题2：输入输出维度不匹配

错误信息：

RuntimeError: Expected input[0] to have 4 dimension(s). Got 3

解决方案：

# 确保输入数据的维度与模型期望一致
# 检查模型输入形状
session = ort.InferenceSession("model.onnx")
input_shape = session.get_inputs()[0].shape
print(f"Expected input shape: {input_shape}")

# 调整输入数据维度
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32)  # 添加batch维度
问题3：动态图转换问题

错误信息：

RuntimeError: ONNX export failed: Could not export a Python function

解决方案：

# 1. 使用torch.jit.script先转换为静态图
model_scripted = torch.jit.script(model)
torch.onnx.export(model_scripted, ...)

# 2. 或者使用torch.jit.trace
model_traced = torch.jit.trace(model, dummy_input)
torch.onnx.export(model_traced, ...)

# 3. 避免在forward中使用Python控制流，改用torch函数
# 错误：if x > 0: return x
# 正确：return torch.where(x > 0, x, 0)

6.2 推理异常类问题

问题1：结果精度差异大

原因分析：

  • 数值精度问题（FP32 vs FP16）
  • 预处理不一致
  • 算子实现差异

解决方案：

# 1. 对比原始框架和ONNX Runtime的输出
import torch
import onnxruntime as ort
import numpy as np

# PyTorch输出
model_pytorch.eval()
with torch.no_grad():
    output_pytorch = model_pytorch(torch.from_numpy(input_data)).numpy()

# ONNX Runtime输出
session = ort.InferenceSession("model.onnx")
output_onnx = session.run(None, {input_name: input_data})[0]

# 计算差异
diff = np.abs(output_pytorch - output_onnx)
print(f"Max difference: {np.max(diff)}")
print(f"Mean difference: {np.mean(diff)}")

# 2. 确保预处理一致
# 3. 使用FP32精度（避免量化导致的精度损失）
问题2：推理速度未达标

优化方案：

# 1. 使用GPU加速
session = ort.InferenceSession(
    "model.onnx",
    providers=['CUDAExecutionProvider']
)

# 2. 启用图优化
options = ort.SessionOptions()
options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
session = ort.InferenceSession("model.onnx", sess_options=options)

# 3. 使用量化模型
# 参考第四章的量化方法

# 4. 调整线程数
options.intra_op_num_threads = 8
options.inter_op_num_threads = 8
问题3：内存占用过高

优化方案：

# 1. 使用量化模型减少内存占用
# 2. 限制batch size
# 3. 使用内存映射（对于大模型）
options = ort.SessionOptions()
options.enable_mem_pattern = False  # 禁用内存模式优化（可能增加内存但减少碎片）

6.3 部署问题类

问题1：平台兼容性

解决方案：

# 1. 检查ONNX Runtime版本和平台支持
import onnxruntime as ort
print("Available providers:", ort.get_available_providers())

# 2. 使用opset版本兼容性
# 较老的opset版本（如opset 9）兼容性更好
torch.onnx.export(..., opset_version=9)

# 3. 测试不同平台的兼容性
问题2：依赖库冲突

解决方案：

# 使用虚拟环境隔离依赖
python -m venv onnx_env
source onnx_env/bin/activate  # Linux/macOS
# onnx_env\Scripts\activate  # Windows

# 安装特定版本
pip install onnx==1.12.0
pip install onnxruntime==1.12.0

第七章：ONNX生态与进阶方向

7.1 ONNX生态周边工具

7.1.1 模型可视化工具
  • Netron：最流行的ONNX模型可视化工具（Web版和桌面版）
  • ONNX GraphSurgeon：Python库，用于修改ONNX模型结构
  • onnx2pytorch：将ONNX模型转换回PyTorch（用于调试）
7.1.2 模型转换工具
  • onnx-tensorflow：ONNX转TensorFlow
  • onnx2keras：ONNX转Keras
  • onnx2torch：ONNX转PyTorch
7.1.3 性能分析工具
  • ONNX Runtime Profiler：分析模型推理性能瓶颈
  • onnxruntime-extensions：扩展ONNX Runtime的自定义算子支持

7.2 进阶学习方向

7.2.1 自定义ONNX算子

场景：当模型包含ONNX不支持的自定义算子时，需要实现自定义算子。

实现步骤：

  1. 定义算子规范：

    # 在ONNX中定义自定义算子
from onnx import helper

custom_op = helper.make_node(
    'CustomOp',  # 算子名称
    ['input'],   # 输入
    ['output'],  # 输出
    domain='custom'  # 自定义域
)

  2. 在ONNX Runtime中实现算子：

    • 使用onnxruntime-extensions
    • 或实现C++扩展
7.2.2 ONNX模型压缩与蒸馏结合

模型压缩流程：

  1. 训练大模型（教师模型）
  2. 知识蒸馏训练小模型（学生模型）
  3. 转换为ONNX格式
  4. 进一步量化优化

示例：

# 1. 训练教师模型（PyTorch）
teacher_model = train_teacher_model()

# 2. 知识蒸馏训练学生模型
student_model = distill_student_model(teacher_model)

# 3. 转换为ONNX
torch.onnx.export(student_model, ...)

# 4. 量化优化
quantize_static(...)
7.2.3 大模型的ONNX转换与部署

挑战：

  • 模型体积大（几GB到几十GB）
  • 内存占用高
  • 推理速度慢

解决方案：

  1. 模型分片：

    # 将大模型拆分为多个子图
# 使用ONNX的subgraph功能

  2. 动态量化：

    # 对大模型进行INT8量化
quantize_dynamic(...)

  3. 使用专门的推理引擎：

    • TensorRT（NVIDIA GPU）
    • OpenVINO（Intel硬件）
    • ONNX Runtime with optimizations

  4. 模型并行：

    • 将模型分布到多个GPU/设备上

总结与展望

核心知识点总结

通过本文的学习，我们掌握了：

  1. ONNX核心概念：

    • ONNX是开放的模型交换格式，实现跨框架、跨平台部署
    • ONNX Runtime是高性能推理引擎，支持多种硬件平台

  2. 模型转换流程：

    • PyTorch：使用torch.onnx.export()
    • TensorFlow/Keras：使用tf2onnx
    • 注意动态维度设置和算子兼容性

  3. 模型优化方法：

    • 图优化：使用onnxoptimizer和onnx-simplifier
    • 量化：动态量化和静态量化，可减少75%模型体积，提升2-4倍速度

  4. 多平台部署：

    • CPU/GPU：通过执行提供者选择硬件
    • 移动端：ONNX Runtime Mobile
    • 边缘设备：结合TensorRT、OpenVINO

  5. 问题排查：

    • 转换失败：检查算子支持、维度匹配、动态图问题
    • 推理异常：对比精度、优化速度、降低内存占用

ONNX发展趋势展望

  1. 支持更多算子：

    • ONNX Operator Set不断更新，支持更多新算子
    • 更好的动态图支持

  2. 大模型部署优化：

    • 针对Transformer架构的优化
    • 支持模型并行和分片

  3. 硬件加速：

    • 更多硬件厂商的支持（如Apple Neural Engine、Google TPU）
    • 更高效的量化方案

  4. 工具链完善：

    • 更易用的转换工具
    • 更强大的可视化工具
    • 更完善的性能分析工具

  5. 生态融合：

    • 与MLOps平台的深度集成
    • 与边缘计算框架的结合

附录

A. 常用学习资源

官方文档
开源项目
优质教程

B. 版本依赖说明

本文代码基于以下版本测试：

Python: 3.8+
PyTorch: 1.12.0+
TensorFlow: 2.10.0+
ONNX: 1.12.0+
ONNX Runtime: 1.12.0+
onnxoptimizer: 0.3.12+
tf2onnx: 1.13.0+

致谢：感谢ONNX社区和所有贡献者的努力，让模型部署变得更加简单高效。

本文持续更新中，如有问题或建议，欢迎反馈。

avatar
文章
阅读
粉丝