ONNX 基础入门
ONNX 是什么
ONNX,即 Open Neural Network Exchange(开放神经网络交换) ,是一种用于表示深度学习模型的开放标准文件格式。它由 Facebook 和 Microsoft 在 2017 年联合开发,后来得到了 NVIDIA、Intel、AWS、Google、OpenAI 等众多公司的支持,旨在解决不同深度学习框架之间模型格式不兼容的问题,为模型的存储、交换和部署提供统一标准,使得不同的机器学习框架(如 PyTorch、TensorFlow、Caffe 等)可以采用相同格式存储模型数据并进行交互。
ONNX 的核心优势在于它定义了一组和环境、平台均无关的标准格式。这意味着无论你使用何种训练框架训练模型,在训练完毕后都可以将这些框架的模型统一转换为 ONNX 这种统一的格式进行存储。ONNX 文件不仅仅存储了神经网络模型的权重,同时也存储了模型的结构信息以及网络中每一层的输入输出和一些其它的辅助信息 。这就好比为不同语言的人提供了一种通用的交流语言,使得模型在不同框架和平台之间的迁移变得更加顺畅。在获得 ONNX 模型之后,模型部署人员自然就可以将这个模型部署到兼容 ONNX 的运行环境中去。
从技术实现角度来看,ONNX 通过 Protobuf 进行序列化,使用 protobuf 子集,既支持 v2 也支持 v3 。它提供了单一文件格式,每个机器学习库都有自己的文件格式,如 PyTorch 可以保存为 '.pth',而 ONNX 为保存和导出模型文件提供了单一标准,文件扩展名为 '.onnx'。ONNX 定义了一个可扩展的计算图模型,为各种框架不同的图形表示提供了标准的图形表示,ONNX 图通过各种计算节点表示模型图,可以使用 Netron 等工具进行可视化;定义了标准数据类型,为图中节点提供通用的数据类型规范,包括 int、float 等;还定义了内置运算符,负责将 ONNX 中的运算符类型映射到所需的框架,比如将 PyTorch 模型转换为 ONNX 时,PyTorch 运算符会映射到 ONNX 中的关联运算符。
ONNX 的应用场景
ONNX 凭借其出色的跨框架和跨平台特性,在多个领域都有着广泛的应用,成为推动深度学习发展的重要力量。
- 计算机视觉领域:ONNX 模型被广泛应用于目标检测、图像分类、语义分割等任务。在目标检测中,像经典的 YOLO 系列模型,开发者可以先使用 PyTorch 或 TensorFlow 等框架进行训练,然后将训练好的模型转换为 ONNX 格式,这样就能方便地部署到不同的硬件平台和推理引擎上。利用 ONNX Runtime 在 NVIDIA GPU 上进行加速推理,或者使用 OpenVINO 在 Intel CPU 上实现高效运行,从而满足不同场景下对检测速度和精度的要求 。在图像分类任务中,无论是基于 ResNet、VGG 等传统架构,还是如 EfficientNet 等新型高效架构的模型,都可以借助 ONNX 实现快速部署,实现对大量图像数据的快速分类。在语义分割方面,例如用于医学影像分析的 U-Net 模型,通过转换为 ONNX 格式,可以在不同的医疗设备或计算平台上运行,帮助医生更准确地识别病变区域。
- 自然语言处理领域:ONNX 也发挥着重要作用,常见于文本分类、情感分析、自动问答等任务。以文本分类为例,使用 Transformer 架构进行训练的模型,转换为 ONNX 格式后,可以在不同的服务器或移动设备上高效运行,实现对用户输入文本的实时分类。在情感分析中,基于预训练语言模型(如 BERT、GPT 等)微调后的模型,通过 ONNX 可以方便地部署到各种应用场景中,分析社交媒体评论、客户反馈等文本的情感倾向。自动问答系统中,ONNX 使得模型能够在不同的硬件环境下快速响应,为用户提供准确的答案 。
- 其他领域:在数据分析和预测领域,如时间序列分析、异常检测、推荐系统等,ONNX 同样适用。对于时间序列分析模型,将其转换为 ONNX 格式后,可以在不同的数据分析平台上运行,对金融数据、气象数据等进行实时预测和分析。在推荐系统中,ONNX 模型可以在电商平台的服务器上高效运行,根据用户的行为数据和商品信息,为用户提供个性化的推荐服务 。在机器人技术领域,ONNX 模型用于机器人视觉、控制和导航等方面。例如,机器人通过视觉摄像头获取图像信息,利用基于 ONNX 的目标检测模型识别周围环境中的物体,从而更加智能地进行操作和导航 。
与其他模型格式的比较
在深度学习中,不同的框架和应用场景衍生出了多种模型格式,ONNX 与其他常见模型格式(如 TensorFlow 的.pb、PyTorch 的.pth)相比,有着各自的特点和优劣。
- PyTorch 的.pth 格式:.pth 文件通常用于保存 PyTorch 模型的权重(parameters)和训练过程中的优化器状态(optimizer state) 。模型状态字典是.pth 文件的主要部分,包含了模型的所有权重,如卷积层的滤波器、全连接层的权重和偏置等,以 Python 字典对象存储,键是层的名字,值是相应的参数张量。当保存模型时,也会保存优化器的状态,便于在模型训练中断后从上次停止的地方继续训练 ,优化器状态字典包含每个参数的梯度、动量和其他与优化器相关的状态信息。.pth 文件有时还会包含其他辅助信息,如训练的损失值、额外的元数据或者批处理大小等,但这些不是必须的。它的优势在于与 PyTorch 框架紧密结合,保存和加载操作简单方便,非常适合在 PyTorch 环境下进行模型的训练和调试 。由于其保存了优化器状态,在继续训练场景下有天然优势。不过,.pth 格式的模型具有框架依赖性,只能在 PyTorch 环境中使用,在跨框架部署时存在困难。
- TensorFlow 的.pb 格式:.pb 文件(Protocol Buffer)是 TensorFlow 的一种模型保存格式,它是一种二进制文件,用于存储模型的图结构(包括节点和边)以及模型的权重参数 。.pb 文件将模型的计算图和参数序列化,使得模型可以方便地在不同环境中部署和运行。在 TensorFlow Serving 中,.pb 文件是常用的模型格式,能够高效地为线上服务提供推理支持。它的优点是适合 TensorFlow 的生产部署,具有较高的运行效率和稳定性 ,可以很好地利用 TensorFlow 的优化和加速技术。缺点是同样存在框架局限性,主要适用于 TensorFlow 生态系统,与其他框架的交互性较差,并且.pb 文件对于模型结构的可视化和调试相对困难。
- ONNX 格式:ONNX 作为一种开放标准格式,最大的优势在于其跨框架兼容性。它可以将不同框架(如 PyTorch、TensorFlow、MXNet 等)训练的模型统一转换为 ONNX 格式,实现模型在不同框架和平台之间的无缝迁移 。ONNX 为模型提供了一个中立的表示形式,使得硬件和软件厂商可以基于 ONNX 标准来优化模型性能,所有支持 ONNX 标准的框架都能从中受益,例如通过 ONNX Runtime 可以大幅提升推理速度,并且支持多种硬件加速,如 NVIDIA TensorRT(GPU 加速)、Intel OpenVINO(CPU/FPGA 加速)、DirectML(Windows 端 GPU 加速) 。ONNX 还便于模型的可视化和分析,使用 Netron 等工具可以直观地查看模型的结构。不过,ONNX 也存在一些不足,它并非支持所有模型层,某些复杂算子(如特定 Transformer 变体)可能不支持,需要手动转换;在将某些 PyTorch/TensorFlow 算子转换时,可能需要手动调整模型结构 。
ONNX 推理部署基础例程
模型导出
在将模型用于推理部署之前,首先需要将其导出为 ONNX 格式。以 PyTorch 框架为例,使用torch.onnx.export函数来实现这一转换过程。下面通过一个简单的线性回归模型来详细说明导出步骤及参数含义。
首先,定义一个简单的线性回归模型:
import torch
import torch.nn as nn
class LinearRegression(nn.Module):
def __init__(self):
super(LinearRegression, self).__init__()
self.linear = nn.Linear(1, 1) # 定义一个线性层,输入和输出维度均为1
def forward(self, x):
out = self.linear(x)
return out
# 初始化模型
model = LinearRegression()
接下来,创建一个虚拟输入,用于指定模型的输入形状。这个虚拟输入的形状需要与模型实际运行时的输入形状一致:
# 创建虚拟输入
dummy_input = torch.randn(1, 1) # 这里生成一个形状为(1, 1)的随机张量作为输入
然后,使用torch.onnx.export函数将模型导出为 ONNX 格式:
# 导出模型为ONNX格式
onnx_path = 'linear_regression.onnx' # 指定导出的ONNX模型文件路径
torch.onnx.export(model, # 要导出的模型
dummy_input, # 用于指定模型输入的虚拟数据,确保模型输入格式正确
onnx_path, # 导出的ONNX模型文件路径
input_names=['input'], # 指定输入的名称为'input'
output_names=['output'], # 指定输出的名称为'output'
export_params=True, # 是否导出模型参数,True表示导出,False则不导出
opset_version=11 # 指定ONNX算子集版本,不同版本支持的算子可能不同
)
在上述代码中:
- model:是要导出的 PyTorch 模型实例。
- dummy_input:是一个虚拟输入张量,通过运行模型的前向传播,来确定模型的计算图结构和输入输出形状。它的形状和数据类型需要与模型实际运行时的输入一致。
- onnx_path:指定了导出的 ONNX 模型文件的保存路径及文件名。
- input_names:是一个列表,用于指定模型输入的名称,在 ONNX 模型中,输入数据将通过这些名称进行标识。
- output_names:同样是一个列表,用于指定模型输出的名称,方便在推理时获取输出结果。
- export_params:设置为True时,表示将模型的参数(如权重和偏置)一并导出到 ONNX 文件中;如果设置为False,则只导出模型的结构,不包含参数。
- opset_version:指定了 ONNX 算子集的版本。不同版本的算子集支持的算子有所不同,在导出模型时,需要根据模型中使用的算子来选择合适的版本。一般建议使用较新的版本,以确保对更多算子的支持 。
模型验证
导出 ONNX 模型后,需要验证其正确性,以确保模型在推理部署过程中能够正常运行。可以使用 ONNX Runtime 来加载并运行导出的 ONNX 模型,并将其输出与原始 PyTorch 模型的输出进行对比。
首先,安装 ONNX Runtime 库:
pip install onnxruntime
然后,编写验证代码:
import onnxruntime
import numpy as np
import torch
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession('linear_regression.onnx')
# 定义输入数据(与导出模型时的虚拟输入形状一致)
input_data = torch.randn(1, 1).numpy()
# ONNX模型推理
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outs = ort_session.run(None, ort_inputs)
# PyTorch模型推理
model = LinearRegression()
model.eval()
with torch.no_grad():
torch_out = model(torch.from_numpy(input_data)).numpy()
# 验证输出结果是否一致
np.testing.assert_allclose(ort_outs[0], torch_out, rtol=1e-03, atol=1e-05)
print("ONNX模型验证通过,输出结果与PyTorch模型一致!")
在上述代码中:
- 使用onnxruntime.InferenceSession加载导出的 ONNX 模型,创建一个推理会话。
- 准备输入数据,这里使用与导出模型时相同形状的随机数据,并将其转换为 NumPy 数组格式。
- 通过推理会话的run方法执行 ONNX 模型的推理,传入输入数据并获取输出结果。
- 在 PyTorch 模型中进行同样的推理操作,得到 PyTorch 模型的输出。
- 最后,使用np.testing.assert_allclose函数对比 ONNX 模型和 PyTorch 模型的输出结果。rtol和atol分别设置了相对误差和绝对误差的容忍度,只要两个输出结果之间的误差在容忍度范围内,就认为验证通过 。如果验证通过,会打印出 “ONNX 模型验证通过,输出结果与 PyTorch 模型一致!” 的提示信息;否则,会抛出异常并显示具体的误差信息 。
基本推理流程
通过 ONNX Runtime 进行基本推理的过程主要包括加载模型、准备输入数据、执行推理和获取输出结果。下面基于前面导出并验证的线性回归 ONNX 模型,详细介绍推理流程。
import onnxruntime
import numpy as np
# 1. 加载ONNX模型
ort_session = onnxruntime.InferenceSession('linear_regression.onnx')
# 2. 准备输入数据
# 生成与模型输入形状匹配的随机数据
input_data = np.random.randn(1, 1).astype(np.float32)
# 3. 执行推理
ort_inputs = {ort_session.get_inputs()[0].name: input_data}
ort_outs = ort_session.run(None, ort_inputs)
# 4. 获取输出结果
output = ort_outs[0]
print("推理输出结果:", output)
在这个过程中:
- 加载 ONNX 模型:使用onnxruntime.InferenceSession加载之前导出的linear_regression.onnx模型,创建一个推理会话ort_session,这个会话包含了模型的结构和参数信息,为后续的推理操作做准备 。
- 准备输入数据:根据模型的输入要求,生成一个形状为(1, 1)的随机数据,并将其转换为np.float32类型的 NumPy 数组。这里的输入数据形状和数据类型必须与模型导出时使用的虚拟输入一致,否则会导致推理错误。
- 执行推理:构建一个输入字典ort_inputs,将输入数据与模型输入节点的名称进行关联,然后调用推理会话的run方法执行推理。run方法的第一个参数为None,表示获取所有输出节点的结果;第二个参数为输入字典ort_inputs,包含了输入数据。执行推理后,会返回一个包含所有输出结果的列表ort_outs。
- 获取输出结果:从推理结果列表ort_outs中取出第一个元素,即模型的输出结果,并打印出来。在实际应用中,可以根据模型的输出情况,对输出结果进行进一步的处理和分析 。
通过以上步骤,就完成了基于 ONNX 模型的基本推理过程。在实际应用中,可能需要根据具体的业务需求和模型特点,对输入数据进行预处理(如归一化、图像增强等),对输出结果进行后处理(如反归一化、解码、分类等),以满足不同场景下的使用要求。
ONNX 推理部署高级例程
动态维度处理
在实际应用中,输入数据的大小可能会有所变化,例如图像的尺寸可能不同。为了使 ONNX 模型能够适应这种变化,需要在导出模型时设置动态维度。以图像分类模型为例,假设模型的输入为图像,通常希望能够处理不同分辨率的图像。
在 PyTorch 中导出动态维度的 ONNX 模型时,可以通过dynamic_axes参数来指定哪些维度是动态的。下面以一个简单的卷积神经网络(CNN)为例:
import torch
import torch.nn as nn
class SimpleCNN(nn.Module):
def __init__(self):
super(SimpleCNN, self).__init__();
self.conv1 = nn.Conv2d(3, 16, kernel_size=3, padding=1);
self.relu1 = nn.ReLU();
self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2);
self.fc1 = nn.Linear(16 * 112 * 112, 10);
def forward(self, x):
out = self.conv1(x);
out = self.relu1(out);
out = self.pool1(out);
out = out.view(-1, 16 * 112 * 112);
out = self.fc1(out);
return out
# 初始化模型
model = SimpleCNN();
model.eval();
# 创建虚拟输入,这里只是示例,实际使用时可根据需求调整
dummy_input = torch.randn(1, 3, 224, 224);
# 导出为ONNX模型,设置动态维度
onnx_path ='simple_cnn_dynamic.onnx';
dynamic_axes = {
'input': {0: 'batch_size', 2: 'height', 3: 'width'},
'output': {0: 'batch_size'}
};
torch.onnx.export(model,
dummy_input,
onnx_path,
input_names=['input'],
output_names=['output'],
dynamic_axes=dynamic_axes,
opset_version=11);
在上述代码中:
- dynamic_axes是一个字典,用于指定哪些维度是动态的。其中,'input'表示输入张量,{0: 'batch_size', 2: 'height', 3: 'width'}表示输入张量的第 0 维(批次维度)、第 2 维(高度维度)和第 3 维(宽度维度)是动态的,'batch_size'、'height'和'width'是自定义的维度名称,用于在 ONNX 模型中标识这些动态维度;'output'表示输出张量,{0: 'batch_size'}表示输出张量的第 0 维(批次维度)是动态的 。
在使用 ONNX Runtime 进行推理时,同样可以处理动态维度的输入。假设我们有不同尺寸的图像作为输入:
import onnxruntime
import numpy as np
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession('simple_cnn_dynamic.onnx');
# 准备不同尺寸的输入数据
input_data_1 = np.random.randn(1, 3, 224, 224).astype(np.float32);
input_data_2 = np.random.randn(1, 3, 128, 128).astype(np.float32);
# 执行推理
ort_inputs_1 = {ort_session.get_inputs()[0].name: input_data_1};
ort_outs_1 = ort_session.run(None, ort_inputs_1);
ort_inputs_2 = {ort_session.get_inputs()[0].name: input_data_2};
ort_outs_2 = ort_session.run(None, ort_inputs_2);
通过上述方式,ONNX 模型能够处理不同尺寸的输入数据,提高了模型的通用性和灵活性,使其更适合实际应用中的各种场景。
模型优化技巧
为了提升 ONNX 模型的推理性能,可以采用多种优化技巧,其中模型量化和图优化是较为常用的方法。
模型量化:模型量化是一种通过减少模型中数据的表示精度来降低模型存储需求和计算量的技术,从而提升推理速度。常见的量化方式包括 8 位整数量化(INT8)和 16 位半精度浮点数量化(FP16)。以使用 ONNX Runtime 进行 INT8 量化为例:
import onnx
import onnxruntime as ort
from onnxruntime.quantization import quantize_dynamic, QuantType
# 加载原始ONNX模型
model_path = 'original_model.onnx';
model = onnx.load(model_path);
# 进行动态量化,将模型量化为INT8
quantized_model_path = 'quantized_model.onnx';
quantize_dynamic(model_path, quantized_model_path, weight_type=QuantType.QUInt8);
# 使用量化后的模型进行推理
ort_session = ort.InferenceSession(quantized_model_path);
# 准备输入数据并执行推理,这里省略具体代码
在上述代码中,quantize_dynamic函数用于对 ONNX 模型进行动态量化,weight_type=QuantType.QUInt8指定将权重量化为 8 位无符号整数 。量化后的模型在推理时,由于数据表示精度降低,计算量减少,从而可以提高推理速度,尤其在一些对精度要求不是特别高的场景中,量化后的模型能够在几乎不损失太多准确性的前提下,显著提升推理效率 。
图优化:图优化是对 ONNX 模型的计算图进行优化,通过合并节点、消除冗余计算等方式来提高模型的执行效率。ONNX Runtime 在加载模型时会自动进行一些基本的图优化,也可以使用onnxoptimizer库进行更高级的优化。首先安装onnxoptimizer库:
pip install onnxoptimizer
然后使用该库对模型进行优化:
import onnx
from onnxoptimizer import optimize
# 加载原始ONNX模型
model_path = 'original_model.onnx';
model = onnx.load(model_path);
# 进行图优化
optimized_model = optimize(model);
# 保存优化后的模型
optimized_model_path = 'optimized_model.onnx';
onnx.save(optimized_model, optimized_model_path);
# 使用优化后的模型进行推理,这里省略具体推理代码
在这个过程中,optimize函数会对模型进行一系列的优化操作,如常量折叠(将图中的静态变量转换为常量)、死代码消除(去除图中未使用的节点)、算子融合(将多条指令合并为一条,比如将线性层和 ReLU 激活函数合并为一个操作)等 。经过优化后的模型,计算图更加简洁高效,能够在推理时减少计算时间和内存占用,提升整体性能 。
多后端推理
ONNX 模型的强大之处在于它可以在多个不同的后端上进行推理,常见的后端包括 CPU、GPU 和 TensorRT。不同后端在性能和适用场景上存在差异,通过对比可以选择最适合具体应用的后端。
CPU 后端推理:使用 ONNX Runtime 在 CPU 上进行推理是最基本的方式,适用于对计算资源要求不高、部署环境为普通 CPU 的场景。代码示例如下:
import onnxruntime
import numpy as np
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession('model.onnx', providers=['CPUExecutionProvider']);
# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32);
ort_inputs = {ort_session.get_inputs()[0].name: input_data};
# 执行推理
ort_outs = ort_session.run(None, ort_inputs);
在上述代码中,通过providers=['CPUExecutionProvider']指定使用 CPU 后端进行推理。CPU 后端通用性强,但在处理大规模数据和复杂模型时,推理速度可能较慢。
GPU 后端推理:如果系统中安装了 NVIDIA GPU 及相关驱动和库,可以使用 ONNX Runtime 的 GPU 后端来加速推理。首先需要安装onnxruntime-gpu库,然后进行如下操作:
import onnxruntime
import numpy as np
# 加载ONNX模型,使用GPU后端
ort_session = onnxruntime.InferenceSession('model.onnx', providers=['CUDAExecutionProvider']);
# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32);
ort_inputs = {ort_session.get_inputs()[0].name: input_data};
# 执行推理
ort_outs = ort_session.run(None, ort_inputs);
通过providers=['CUDAExecutionProvider']指定使用 CUDA 执行提供程序,利用 GPU 的并行计算能力来加速推理过程。GPU 后端在处理大规模矩阵运算和复杂神经网络时,能够显著提高推理速度,适用于对推理速度要求较高的场景,如图像识别、视频分析等 。
TensorRT 后端推理:TensorRT 是 NVIDIA 推出的高性能深度学习推理优化器和运行时引擎,能够进一步优化 ONNX 模型在 NVIDIA GPU 上的推理性能。使用 TensorRT 进行推理需要先将 ONNX 模型转换为 TensorRT 引擎,示例代码如下:
import tensorrt as trt
import pycuda.driver as cuda
import pycuda.autoinit
import numpy as np
# 创建TensorRT日志记录器
TRT_LOGGER = trt.Logger(trt.Logger.WARNING);
# 加载ONNX模型并创建TensorRT引擎
def build_engine(onnx_model_path):
with trt.Builder(TRT_LOGGER) as builder, builder.create_network(
trt.EXPLICIT_BATCH) as network, trt.OnnxParser(network, TRT_LOGGER) as parser:
builder.max_workspace_size = 1 << 30; # 设置最大工作空间大小
builder.max_batch_size = 1;
with open(onnx_model_path, 'rb') as model:
parser.parse(model.read());
return builder.build_cuda_engine(network);
# 执行TensorRT推理
def infer(engine, input_data):
context = engine.create_execution_context();
# 分配输入和输出缓冲区
input_volume = trt.volume(engine.get_binding_shape(0));
input_batch = input_data.shape[0];
input_size = input_volume * input_batch * np.dtype(np.float32).itemsize;
output_volume = trt.volume(engine.get_binding_shape(1));
output_size = output_volume * input_batch * np.dtype(np.float32).itemsize;
d_input = cuda.mem_alloc(input_size);
d_output = cuda.mem_alloc(output_size);
h_output = np.empty(output_volume * input_batch, dtype=np.float32);
stream = cuda.Stream();
# 将输入数据传输到GPU
cuda.memcpy_htod_async(d_input, input_data, stream);
# 执行推理
context.execute_async_v2(bindings=[int(d_input), int(d_output)], stream_handle=stream.handle);
# 将输出数据从GPU传输回CPU
cuda.memcpy_dtoh_async(h_output, d_output, stream);
stream.synchronize();
return h_output.reshape((input_batch, -1));
# ONNX模型路径
onnx_model_path ='model.onnx';
# 构建TensorRT引擎
engine = build_engine(onnx_model_path);
# 准备输入数据
input_data = np.random.randn(1, 3, 224, 224).astype(np.float32);
# 执行推理
output = infer(engine, input_data);
在上述代码中,build_engine函数负责将 ONNX 模型转换为 TensorRT 引擎,infer函数使用创建好的引擎进行推理。TensorRT 通过对模型进行优化,如层融合、量化等,能够在 GPU 上实现更高的推理性能,特别适用于对推理速度和吞吐量要求极高的生产环境 。
通过对比不同后端的推理性能,可以发现 GPU 后端和 TensorRT 后端在处理复杂模型和大规模数据时,推理速度明显优于 CPU 后端。在实际应用中,应根据具体的硬件环境、模型特点和性能需求,选择最合适的后端进行 ONNX 模型推理 。
实战案例分析
计算机视觉案例
以目标检测模型 YOLOv5 为例,展示从模型导出到 ONNX 推理部署的完整流程。YOLOv5 是一种广泛应用于目标检测的深度学习模型,具有速度快、精度高的特点 。
模型导出:假设已经训练好了 YOLOv5 模型,其权重文件为yolov5s.pt。使用 YOLOv5 官方提供的export.py脚本可以将 PyTorch 模型导出为 ONNX 格式。首先确保安装了必要的依赖库,包括torch、onnx等。
# 安装依赖库
pip install torch onnx
然后执行导出命令:
python export.py --weights yolov5s.pt --include onnx
在上述命令中,--weights指定了要导出的 PyTorch 模型权重文件路径,--include onnx表示导出为 ONNX 格式 。导出过程中,export.py脚本会加载 PyTorch 模型,并根据模型的结构和权重生成对应的 ONNX 模型。
模型验证:导出 ONNX 模型后,需要验证其正确性。可以使用 ONNX Runtime 来加载并运行导出的 ONNX 模型,并将其输出与原始 PyTorch 模型的输出进行对比。首先安装 ONNX Runtime 库:
pip install onnxruntime
然后编写验证代码:
import onnxruntime
import torch
import cv2
import numpy as np
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession('yolov5s.onnx')
# 加载PyTorch模型
model = torch.load('yolov5s.pt', map_location=torch.device('cpu'))['model'].float().eval()
# 准备输入数据,这里以一张测试图片为例
img = cv2.imread('test.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (640, 640))
img = img.transpose((2, 0, 1)).astype(np.float32) / 255.0
img = np.expand_dims(img, axis=0)
# ONNX模型推理
ort_inputs = {ort_session.get_inputs()[0].name: img}
ort_outs = ort_session.run(None, ort_inputs)
# PyTorch模型推理
with torch.no_grad():
torch_out = model(torch.from_numpy(img)).numpy()
# 验证输出结果是否一致
np.testing.assert_allclose(ort_outs[0], torch_out, rtol=1e-03, atol=1e-05)
print("ONNX模型验证通过,输出结果与PyTorch模型一致!")
在上述代码中,首先加载 ONNX 模型和 PyTorch 模型,然后准备一张测试图片作为输入数据。对 ONNX 模型和 PyTorch 模型分别进行推理,最后使用np.testing.assert_allclose函数对比两者的输出结果。如果验证通过,会打印出 “ONNX 模型验证通过,输出结果与 PyTorch 模型一致!” 的提示信息 。
ONNX 推理部署:在实际应用中,使用 ONNX Runtime 进行推理部署。假设部署环境为 Python,示例代码如下:
import onnxruntime
import cv2
import numpy as np
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession('yolov5s.onnx')
# 准备输入数据,这里以一张测试图片为例
img = cv2.imread('test.jpg')
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
img = cv2.resize(img, (640, 640))
img = img.transpose((2, 0, 1)).astype(np.float32) / 255.0
img = np.expand_dims(img, axis=0)
# ONNX模型推理
ort_inputs = {ort_session.get_inputs()[0].name: img}
ort_outs = ort_session.run(None, ort_inputs)
# 后处理,这里简单打印输出结果
output = ort_outs[0]
print("推理输出结果:", output)
在这个过程中,首先加载 ONNX 模型,然后对输入图片进行预处理,使其符合模型的输入要求。通过 ONNX Runtime 执行推理,得到输出结果。在实际应用中,还需要对输出结果进行后处理,如非极大值抑制(NMS)等,以得到最终的目标检测框和类别信息 。
自然语言处理案例
以文本分类模型 BERT 为例,介绍如何将训练好的模型转换为 ONNX 格式并进行推理。BERT(Bidirectional Encoder Representations from Transformers)是一种预训练的语言表示模型,在自然语言处理任务中表现出色 。
模型导出:使用 Hugging Face 的transformers库加载预训练的 BERT 模型,并将其导出为 ONNX 格式。首先确保安装了必要的依赖库,包括torch、transformers、onnx等。
# 安装依赖库
pip install torch transformers onnx
然后编写导出代码:
import torch
from transformers import BertTokenizer, BertForSequenceClassification
# 加载预训练的BERT模型和分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
# 准备输入数据,这里以一个示例句子为例
sentence = "This is a sample sentence for text classification."
encoded_input = tokenizer(sentence, return_tensors='pt')
# 设置ONNX模型的保存路径
onnx_model_path = 'bert_classification.onnx'
# 导出模型为ONNX格式
model.eval()
with torch.no_grad():
torch.onnx.export(model,
(encoded_input['input_ids'], encoded_input['attention_mask']),
onnx_model_path,
input_names=['input_ids', 'attention_mask'],
output_names=['logits'],
opset_version=14,
dynamic_axes={'input_ids': {0: 'batch_size', 1:'sequence_length'},
'attention_mask': {0: 'batch_size', 1:'sequence_length'},
'logits': {0: 'batch_size', 1: 'num_classes'}})
print(f"ONNX模型已导出到 {onnx_model_path}")
在上述代码中,首先加载预训练的 BERT 模型和分词器,然后准备一个示例句子作为输入数据。对句子进行编码,得到input_ids和attention_mask。设置 ONNX 模型的保存路径,并使用torch.onnx.export函数将模型导出为 ONNX 格式。其中,input_names和output_names分别指定了输入和输出的名称,opset_version指定了 ONNX 算子集版本,dynamic_axes指定了动态维度,以便处理不同长度的输入序列 。
模型验证:导出 ONNX 模型后,同样使用 ONNX Runtime 来验证其正确性。安装 ONNX Runtime 库(如果尚未安装):
pip install onnxruntime
然后编写验证代码:
import onnxruntime
import torch
from transformers import BertTokenizer
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession('bert_classification.onnx')
# 加载预训练的BERT分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# 准备输入数据,这里使用与导出模型时相同的示例句子
sentence = "This is a sample sentence for text classification."
encoded_input = tokenizer(sentence, return_tensors='pt')
# ONNX模型推理
ort_inputs = {
'input_ids': encoded_input['input_ids'].numpy(),
'attention_mask': encoded_input['attention_mask'].numpy()
}
ort_outs = ort_session.run(None, ort_inputs)
# PyTorch模型推理
model = torch.load('bert_classification.pth', map_location=torch.device('cpu'))
model.eval()
with torch.no_grad():
torch_out = model(encoded_input['input_ids'], encoded_input['attention_mask']).logits.numpy()
# 验证输出结果是否一致
np.testing.assert_allclose(ort_outs[0], torch_out, rtol=1e-03, atol=1e-05)
print("ONNX模型验证通过,输出结果与PyTorch模型一致!")
在上述代码中,首先加载 ONNX 模型和 BERT 分词器,然后准备输入数据。对 ONNX 模型和 PyTorch 模型分别进行推理,最后对比两者的输出结果,验证 ONNX 模型的正确性 。
ONNX 推理部署:在实际应用中,使用 ONNX Runtime 进行推理部署。假设部署环境为 Python,示例代码如下:
import onnxruntime
import torch
from transformers import BertTokenizer
# 加载ONNX模型
ort_session = onnxruntime.InferenceSession('bert_classification.onnx')
# 加载预训练的BERT分词器
tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
# 准备输入数据,这里以一个新的句子为例
new_sentence = "Another sample sentence for text classification."
encoded_input = tokenizer(new_sentence, return_tensors='pt')
# ONNX模型推理
ort_inputs = {
'input_ids': encoded_input['input_ids'].numpy(),
'attention_mask': encoded_input['attention_mask'].numpy()
}
ort_outs = ort_session.run(None, ort_inputs)
# 后处理,这里简单打印输出结果
logits = ort_outs[0]
predicted_class = np.argmax(logits, axis=1)
print("预测类别:", predicted_class)
在这个过程中,首先加载 ONNX 模型和 BERT 分词器,然后对新的句子进行编码并准备输入数据。通过 ONNX Runtime 执行推理,得到输出结果。对输出结果进行简单的后处理,如通过np.argmax函数获取预测的类别,并打印出来。在实际应用中,还可以根据具体需求对预测结果进行更复杂的处理,如输出类别名称、计算置信度等 。
常见问题与解决方案
模型转换问题
在将深度学习模型转换为 ONNX 格式的过程中,常常会遭遇各种问题,这些问题会阻碍模型的顺利转换与后续的推理部署,以下是一些常见问题及对应的解决方法:
- 算子不支持:由于不同深度学习框架和 ONNX 标准之间的差异,某些算子在 ONNX 中可能不受支持。例如,一些 PyTorch 中的自定义算子或特定版本的算子,在转换为 ONNX 时会遇到困难。当 PyTorch 模型中使用了特定的自定义激活函数,而 ONNX 中没有对应的标准算子,就会导致转换失败 。解决这个问题,可以尝试更新 PyTorch 和 ONNX 的版本,因为新版本往往会增加对更多算子的支持;也可以使用第三方库,如 ONNX-TensorRT、ONNXRuntime 等,这些库提供额外功能和优化,增加对不支持算子的支持;还可以进行自定义实现,如果对模型的具体算子实现非常了解,可尝试手动实现这些不支持的算子,并将其添加到转换过程中 。
- 数据类型不匹配:深度学习框架内部的数据类型表示方式与 ONNX 的标准数据类型可能存在不一致的情况。在 PyTorch 中使用了双精度浮点数(float64),而 ONNX 默认支持的是单精度浮点数(float32),转换时就会出现数据类型不匹配的错误 。为解决数据类型不匹配问题,在模型训练阶段,尽量使用 ONNX 支持的数据类型;在转换模型前,仔细检查模型中各层的数据类型,并进行必要的转换;如果模型必须使用特定的数据类型,可以在转换后,通过自定义代码对数据类型进行调整 。
- 动态图相关问题:PyTorch 采用动态图机制,模型的结构在每次迭代中都可以变化,而 ONNX 使用静态图,这会导致在导出过程中出现一致性问题。在动态图中,模型的某些分支或循环结构可能难以直接映射到静态图中 。为了解决动态图相关问题,在导出模型前,确保模型的结构是固定的,避免使用动态生成的层或结构;对于复杂的控制流结构,可以尝试将其转换为等价的静态结构,或者使用 ONNX 支持的控制流算子(如果有)来表示 。
推理性能问题
在 ONNX 模型的推理过程中,可能会出现推理速度慢、内存占用高等性能问题,这些问题会影响模型在实际应用中的效果,以下是对这些问题的分析及优化建议:
- 推理速度慢:推理速度慢是 ONNX 模型推理中常见的问题,导致这一问题的原因多种多样。如果模型未使用 INT8 或 FP16 精度优化,会导致计算量较大,从而使推理速度变慢;若 ONNX Runtime 后端配置不当,未启用适合目标硬件的执行提供程序(如 CUDA、TensorRT),也会影响推理速度;输入分辨率过高,会增加计算负担,进而降低推理速度;批量大小设置不合理,过小或过大均可能导致硬件资源未被充分利用,影响推理速度 。针对这些问题,可以采取模型量化的方法,将模型从 FP32 转换为 INT8 或 FP16,减少计算量;根据硬件类型(如 GPU 或 CPU),启用适配的后端(如 CUDA、TensorRT);在满足精度要求的前提下,降低输入分辨率以减少计算负担;测试不同批量大小,找到最适合硬件资源配置的值 。
- 内存占用高:深度学习模型本身参数众多,计算复杂,在推理过程中需要存储大量的中间结果和模型参数,这会占用大量内存。当处理高分辨率图像或长序列数据时,内存需求会进一步增加。此外,不合理的模型结构和推理算法也会导致内存占用过高 。为降低内存的使用,可采用模型剪枝技术,去除模型中不重要的连接和节点,减少模型参数数量,从而降低内存占用;进行内存优化,在推理过程中,合理管理内存,及时释放不再使用的中间结果;使用内存共享技术,多个计算任务之间共享内存,减少内存的重复分配 。
总结与展望
回顾主要内容
本文深入探讨了 ONNX 推理部署的相关知识与实践。开篇介绍 ONNX 作为一种开放标准文件格式,旨在解决深度学习框架间模型格式不兼容问题,凭借其跨框架和跨平台特性,在计算机视觉、自然语言处理等多领域广泛应用,相比其他模型格式,ONNX 具有良好的跨框架兼容性,但也存在对复杂算子支持不足等问题 。
在推理部署基础例程部分,以 PyTorch 框架下的线性回归模型为例,详细阐述了模型导出为 ONNX 格式的步骤,包括使用torch.onnx.export函数时各参数的含义;利用 ONNX Runtime 加载模型并与原始 PyTorch 模型对比输出结果,验证 ONNX 模型的正确性;以及通过 ONNX Runtime 实现基本推理,涵盖加载模型、准备输入数据、执行推理和获取输出结果等关键步骤 。
高级例程中,为使 ONNX 模型适应输入数据大小变化,通过dynamic_axes参数设置动态维度,实现对不同尺寸输入的处理;采用模型量化(如 INT8 量化)和图优化(利用onnxoptimizer库)等技巧提升模型推理性能;对比 CPU、GPU 和 TensorRT 等多后端推理,展示了不同后端在性能和适用场景上的差异,开发者可根据实际需求选择合适后端 。
在实战案例分析环节,以目标检测模型 YOLOv5 和文本分类模型 BERT 为例,从模型导出、验证到推理部署,展示了 ONNX 在计算机视觉和自然语言处理领域的完整应用流程 。最后,针对模型转换和推理性能方面的常见问题,如算子不支持、数据类型不匹配、推理速度慢、内存占用高等,给出了相应的解决方案和优化建议 。
ONNX 未来发展趋势
展望未来,ONNX 在深度学习领域将持续发挥重要作用并不断演进。在对新硬件的支持方面,随着硬件技术的飞速发展,新的计算芯片和架构不断涌现,如 Google 的 TPU、华为的昇腾芯片等。ONNX 有望进一步拓展对这些新型硬件的支持,通过与硬件厂商的紧密合作,优化模型在新硬件上的执行效率,充分发挥硬件的计算能力 。对于边缘计算设备,如物联网传感器、智能摄像头、移动终端等,ONNX 将不断适配其有限的计算资源和存储容量,提供更高效的模型部署方案,推动人工智能在边缘设备上的广泛应用 。
在优化算法上,ONNX 将不断探索和引入更先进的优化算法。在模型量化方面,除了现有的 INT8 和 FP16 量化,可能会发展出更细粒度、更高精度的量化方法,在进一步降低模型计算量和存储需求的同时,最大程度减少对模型精度的影响 。在图优化方面,将不断改进和创新优化策略,如更智能的算子融合算法,能够更精准地识别和合并计算图中的相关算子,进一步减少计算步骤和内存占用,提高模型的推理速度 。同时,随着人工智能技术的不断发展,ONNX 可能会与新兴技术如量子计算、神经形态计算等相结合,为深度学习模型的推理和训练带来全新的思路和方法,拓展 ONNX 的应用边界 。
ONNX 作为深度学习模型部署的关键技术,未来将在新硬件支持和优化算法等方面持续创新和发展,为人工智能技术的广泛应用和性能提升提供有力支撑 。