ajax网站开发技术湖南建设人才网

【深度学习】【OnnxRuntime】【Python】模型转化、环境搭建以及模型部署的详细教程

提示:博主取舍了很多大佬的博文并亲测有效,分享笔记邀大家共同学习讨论

---

前言

ONNXRuntime是微软推出的一款高性能的机器学习推理引擎框架，用户可以非常便利的用其运行一个onnx模型，专注于加速机器学习模型的预测阶段。ONNXRuntime设计目的是为了提供一个高效的执行环境，使机器学习模型能够在各种硬件上快速执行，支持多种运行后端包括CPU，GPU，TensorRT，DML等，使得开发者可以灵活选择最适合其应用场景的硬件平台。
ONNXRuntime是对ONNX模型最原生的支持。

读者可以通过学习【onnx部署】部署系列学习文章目录的onnxruntime系统学习–Python篇 的内容，系统的学习OnnxRuntime部署不同任务的onnx模型。

---

模型转换–pytorch转onnx

Pytorch模型转onnx并推理的步骤如下：

1. 将PyTorch预训练模型文件( .pth 或 .pt 格式)转换成ONNX格式的文件(.onnx格式)，这一转换过程在PyTorch环境中进行。
2. 将转换得到的 .onnx 文件随后作为输入，调用ONNXRuntime的C++ API来执行模型的推理。

博主使用AlexNet图像分类(五种花分类)进行演示，需要安装pytorch环境，对于该算法的基础知识，可以参考博主【AlexNet模型算法Pytorch版本详解】博文

conda create --name AlexNet python==3.10
conda activate AlexNet
# 根据自己主机配置环境
pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu118
# 假设模型转化出错则降级为指定1.16.1版本
pip install onnx==1.16.1

然后把训练模型好的AlexNet.pth模型转成AlexNet.onnx模型，pyorch2onnx.py转换代码如下：

import torch
from model import AlexNet
model = AlexNet(num_classes=5)
weights_path = "./AlexNet.pth"
# 加载模型权重
model.load_state_dict(torch.load(weights_path))
# 模型推理模式
model.eval()
model.cpu()
# 虚拟输入数据
dummy_input1 = torch.randn(1, 3, 224, 224)
# 模型转化函数
torch.onnx.export(model, (dummy_input1), "AlexNet.onnx", verbose=True, opset_version=11)

【AlexNet.pth百度云链接，提取码：ktq5 】直接下载使用即可。

---

Windows平台搭建依赖环境

需要在anaconda虚拟环境安装onnxruntime，需要注意onnxruntime-gpu, cuda, cudnn三者的版本要对应，具体参照官方说明。

博主是win11+cuda12.1+cudnn8.8.1，对应onnxruntime-gpu==1.18.0

import torch
# 查询cuda版本
print(torch.version.cuda)
# 查询cudnn版本
print(torch.backends.cudnn.version())

# 激活环境
activate AlexNet
# 安装onnx
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnx
# 安装GPU版
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnxruntime-gpu==1.18.0
# 或者可以安装CPU版本:没有版本对应要求
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple onnxruntime
# 安装opencv
pip install -i https://pypi.tuna.tsinghua.edu.cn/simple opencv-python

---

onnxruntime调用onnx模型

ONNXRuntime推理核心流程

设置会话选项
通常包括配置优化器级别、线程数和设备(GPU/CPU)使用等。

sess_options = ort.SessionOptions()
sess_options.log_severity_level = 3
sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
sess_options.intra_op_num_threads = 4
providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']

会换选项	日志严重性级别	优化器级别	线程数	设备使用
选项	log_severity_level	graph_optimization_level	graph_optimization_level	CUDAExecutionProvider;CPUExecutionProvider
作用	决定了哪些级别的日志信息将被记录下来，运行时提供了几个预定义的宏来表示不同的日志级别。	在模型加载到ONNXRuntime之前对其进行图优化的过程，提高执行效率	设置每个运算符内部执行时的最大线程数	CUDA/CPU设备选择。
参数	整形，1:Info, 2:Warning. 3:Error, 4:Fatal，默认是2。	ORT_ENABLE_BASIC:基本的图优化； ORT_DISABLE_ALL:禁用所有优化；ORT_ENABLE_EXTENDED:启用扩展优化；ORT_ENABLE_ALL:启用所有优化。	整型	列表中的顺序决定了执行提供者的优先级。

加载模型并创建会话
加载预训练的ONNX模型文件，使用运行时环境、会话选项和模型创建一个Session对象。

session = ort.InferenceSession(onnxpath, sess_options=sess_options, providers=providers)

ort.InferenceSession参数	path_or_bytes	sess_options	providers
内容	模型的位置或者模型的二进制数据	会话选项	设备选择

获取模型输入输出信息
从Session对象中获取模型输入和输出的详细信息，包括数量、名称、类型和形状。

input_nodes_num = len(session.get_inputs()) 
output_nodes_num = len(session.get_outputs()) 
input_name = session.get_inputs()[i].name
output_name = session.get_outputs()[i].name
input_shape = session.get_inputs()[i].shape
output_shape = session.get_outputs()[i].shape

预处理输入数据
对输入数据进行颜色空间转换，尺寸缩放、标准化以及形状维度扩展操作。

rgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
blob = cv2.resize(rgb, (input_w, input_h))
blob = blob.astype(np.float32)
blob /= 255.0
blob -= np.array([0.485, 0.456, 0.406])
blob /= np.array([0.229, 0.224, 0.225])
timg = cv2.dnn.blobFromImage(blob)

这部分不是OnnxRuntime核心部分，根据任务需求不同，代码略微不同。

执行推理
调用Session.run方法，传入输入张量、输出张量名和其他必要的参数，执行推理。

ort_outputs = session.run(output_names=input_node_names, input_feed={output_node_names[0]: timg})

Session.run参数	output_names	input_feed
含义	输出节点名称的列表。	输入节点名称和输入数据的键值对字典，可能有多个输入。

后处理推理结果
推理完成后，从输出张量中获取结果数据，根据需要对结果进行后处理，以获得最终的预测结果。

prob = ort_outputs[0]
max_index = np.argmax(prob)

这部分不是OnnxRuntime核心部分，根据任务需求不同，代码基本不同。

---

ONNXRuntime推理代码

需要配置flower_classes.txt文件存储五种花的分类标签，并将其放置到工程目录下(推荐)。

daisy
dandelion
roses
sunflowers
tulips

这里需要将AlexNet.onnx放置到工程目录下(推荐)，并且将以下推理代码拷贝到新建的py文件中，并执行查看结果。

import onnxruntime as ort
import cv2
import numpy as np# 加载标签文件获得分类标签
def read_class_names(file_path="./flower_classes.txt"):class_names = []try:with open(file_path, 'r') as fp:for line in fp:name = line.strip()if name:class_names.append(name)except IOError:print("could not open file...")import syssys.exit(-1)return class_names# 主函数
def main():# 预测的目标标签数labels = read_class_names()# 测试图片image_path = "./sunflowers.jpg"image = cv2.imread(image_path)cv2.imshow("输入图", image)cv2.waitKey(0)# 设置会话选项sess_options = ort.SessionOptions()# 0=VERBOSE, 1=INFO, 2=WARN, 3=ERROR, 4=FATALsess_options.log_severity_level = 3# 优化器级别:基本的图优化级别sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC# 线程数:4sess_options.intra_op_num_threads = 4# 设备使用优先使用GPU而是才是CPU,列表中的顺序决定了执行提供者的优先级providers = ['CUDAExecutionProvider', 'CPUExecutionProvider']# onnx训练模型文件onnxpath = "./AlexNet.onnx"# 加载模型并创建会话session = ort.InferenceSession(onnxpath, sess_options=sess_options, providers=providers)input_nodes_num = len(session.get_inputs())     # 输入节点输output_nodes_num = len(session.get_outputs())   # 输出节点数input_node_names = []                           # 输入节点名称output_node_names = []                          # 输出节点名称# 获取模型输入信息for i in range(input_nodes_num):# 获得输入节点的名称并存储input_name = session.get_inputs()[i].nameinput_node_names.append(input_name)# 显示输入图像的形状input_shape = session.get_inputs()[i].shapech, input_h, input_w = input_shape[1], input_shape[2], input_shape[3]print(f"input format: {ch}x{input_h}x{input_w}")# 获取模型输出信息for i in range(output_nodes_num):# 获得输出节点的名称并存储output_name = session.get_outputs()[i].nameoutput_node_names.append(output_name)# 显示输出结果的形状output_shape = session.get_outputs()[i].shapenum, nc = output_shape[0], output_shape[1]print(f"output format: {num}x{nc}")input_shape = session.get_inputs()[0].shapeinput_h, input_w = input_shape[2], input_shape[3]print(f"input format: {input_shape[1]}x{input_h}x{input_w}")# 预处理输入数据# 默认是BGR需要转化成RGBrgb = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)# 对图像尺寸进行缩放blob = cv2.resize(rgb, (input_w, input_h))blob = blob.astype(np.float32)# 对图像进行标准化处理blob /= 255.0   # 归一化blob -= np.array([0.485, 0.456, 0.406])  # 减去均值blob /= np.array([0.229, 0.224, 0.225])  # 除以方差#CHW-->NCHW 维度扩展timg = cv2.dnn.blobFromImage(blob)# ---blobFromImage 可以用以下替换---# blob = blob.transpose(2, 0, 1)# blob = np.expand_dims(blob, axis=0)# -------------------------------# 模型推理try:ort_outputs = session.run(output_names=output_node_names, input_feed={input_node_names[0]: timg})except Exception as e:print(e)ort_outputs = None# 后处理推理结果prob = ort_outputs[0]max_index = np.argmax(prob)     # 获得最大值的索引print(f"label id: {max_index}")# 在测试图像上加上预测的分类标签label_text = labels[max_index]cv2.putText(image, label_text, (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.0, (0, 0, 255), 2, 8)cv2.imshow("输入图像", image)cv2.waitKey(0)if __name__ == '__main__':main()

图片正确预测为向日葵：

---

总结

尽可能简单、详细的介绍了pytorch模型到onnx模型的转化，python下onnxruntime环境的搭建以及ONNX模型的OnnxRuntime部署。