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

【深度学习】【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']

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

session = ort.InferenceSession(onnxpath, sess_options=sess_options, providers=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})

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

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 sys
        sys.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=FATAL
    sess_options.log_severity_level = 3
    # 优化器级别:基本的图优化级别
    sess_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
    # 线程数:4
    sess_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].name
        input_node_names.append(input_name)
        # 显示输入图像的形状
        input_shape = session.get_inputs()[i].shape
        ch, 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].name
        output_node_names.append(output_name)
        # 显示输出结果的形状
        output_shape = session.get_outputs()[i].shape
        num, nc = output_shape[0], output_shape[1]
        print(f"output format: {num}x{nc}")

    input_shape = session.get_inputs()[0].shape
    input_h, input_w = input_shape[2], input_shape[3]
    print(f"input format: {input_shape[1]}x{input_h}x{input_w}")

    # 预处理输入数据
    # 默认是BGR需要转化成RGB
    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])  # 除以方差
    #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部署。