PyTorch 模型 onnx 文件导出及调用详情

更新时间：2022年07月25日 14:19:02 作者：荷碧·TongZJ

这篇文章主要介绍了PyTorch模型onnx文件导出及调用详情，文章围绕主题展开详细的内容介绍，具有一定的参考价值，需要的小伙伴可以参考一下

前言

Open Neural Network Exchange (ONNX，开放神经网络交换) 格式，是一个用于表示深度学习模型的标准，可使模型在不同框架之间进行转移

PyTorch 所定义的模型为动态图，其前向传播是由类方法定义和实现的

但是 Python 代码的效率是比较底下的，试想把动态图转化为静态图，模型的推理速度应当有所提升

PyTorch 框架中，torch.onnx.export 可以将父类为 nn.Module 的模型导出到 onnx 文件中，

最重要的有三个参数：

model：父类为 nn.Module 的模型
args：传入 model 的 forward 方法的变量列表，类型应为
tuplef：onnx 文件名称的字符串

import torch
from torchvision.models import resnet50
 
file = 'resnet.onnx'
# 声明模型
resnet = resnet50(pretrained=False).eval()
image = torch.rand([1, 3, 224, 224])
# 导出为 onnx 文件
torch.onnx.export(resnet, (image,), file)

onnx 文件可被 Netron 打开，以查看模型结构

基本用法

要在 Python 中运行 onnx 模型，需要下载 onnxruntime

# 选其一即可
pip install onnxruntime        # CPU 版本
pip install onnxruntime-gpu    # GPU 版本

推理时需要借助其中的 InferenceSession，其中较为重要的实例方法有：

get_inputs()：得到输入变量的列表 (变量属性：name、shape、type)
get_outputs()：得到输入变量的列表 (变量属性：name、shape、type)run(output_names, input_feed)：输入变量为 numpy.ndarray (注意 dtype 应为 float32)，使用模型推理并返回输出

可得出 onnx 模型的基本用法：

import onnxruntime as ort
import numpy as np
file = 'resnet.onnx'
# 找到 GPU / CPU
provider = ort.get_available_providers()[
    1 if ort.get_device() == 'GPU' else 0]
print('设备:', provider)
# 声明 onnx 模型
model = ort.InferenceSession(file, providers=[provider])
# 参考: ort.NodeArg
for node_list in model.get_inputs(), model.get_outputs():
    for node in node_list:
        attr = {'name': node.name,
                'shape': node.shape,
                'type': node.type}
        print(attr)
    print('-' * 60)
 
# 得到输入、输出结点的名称
input_node_name = model.get_inputs()[0].name
ouput_node_name = [node.name for node in model.get_outputs()]
image = np.random.random([1, 3, 224, 224]).astype(np.float32)
print(model.run(output_names=ouput_node_name,
                input_feed={input_node_name: image}))

高级 API

为了简化使用步骤，使用类进行封装：

class Onnx_Module(ort.InferenceSession):
    ''' onnx 推理模型
        provider: 优先使用 GPU'''
    provider = ort.get_available_providers()[
        1 if ort.get_device() == 'GPU' else 0]
 
    def __init__(self, file):
        super(Onnx_Module, self).__init__(file, providers=[self.provider])
        # 参考: ort.NodeArg
        self.inputs = [node_arg.name for node_arg in self.get_inputs()]
        self.outputs = [node_arg.name for node_arg in self.get_outputs()]
 
    def __call__(self, *arrays):
        input_feed = {name: x for name, x in zip(self.inputs, arrays)}
        return self.run(self.outputs, input_feed)

在 PyTorch 中，对于卷积神经网络 model 与图像 image，推理的代码为 "model(image)"，而使用这个封装的类也是类似：

import numpy as np
file = 'resnet.onnx'
model = Onnx_Module(file)
image = np.random.random([1, 3, 224, 224]).astype(np.float32)
print(model(image))

为了方便观察 Torch 模型与 onnx 模型的速度差异，同时检查两个模型的输出是否一致，又编写了 test 函数

test 方法的参数与 torch.onnx.export 一致，其基本流程为：

得到 Torch 模型的输出，并 print 推断耗时
将 Torch 模型导出为 onnx 文件，将输入变量中的 torch.tensor 转化为 numpy.ndarray
初始化 onnx 模型，得到 onnx 模型的输出，并 print 推断耗时
计算 Torch 模型与 onnx 模型输出的绝对误差的均值
将 onnx 模型 return

class Timer:
    repeat = 3
 
    def __new__(cls, fun, *args, **kwargs):
        import time
        start = time.time()
        for _ in range(cls.repeat): fun(*args, **kwargs)
        cost = (time.time() - start) / cls.repeat
        return cost * 1e3  # ms
 
 
class Onnx_Module(ort.InferenceSession):
    ''' onnx 推理模型
        provider: 优先使用 GPU'''
    provider = ort.get_available_providers()[
        1 if ort.get_device() == 'GPU' else 0]
 
    def __init__(self, file):
        super(Onnx_Module, self).__init__(file, providers=[self.provider])
        # 参考: ort.NodeArg
        self.inputs = [node_arg.name for node_arg in self.get_inputs()]
        self.outputs = [node_arg.name for node_arg in self.get_outputs()]
    def __call__(self, *arrays):
        input_feed = {name: x for name, x in zip(self.inputs, arrays)}
        return self.run(self.outputs, input_feed)
 
    @classmethod
    def test(cls, model, args, file, **export_kwargs):
        # 测试 Torch 的运行时间
        torch_output = model(*args).data.numpy()
        print(f'Torch: {Timer(model, *args):.2f} ms')
        # model: Torch -> onnx
        torch.onnx.export(model, args, file, **export_kwargs)
        # data: tensor -> array
        args = tuple(map(lambda tensor: tensor.data.numpy(), args))
        onnx_model = cls(file)
        # 测试 onnx 的运行时间
        onnx_output = onnx_model(*args)
        print(f'Onnx: {Timer(onnx_model, *args):.2f} ms')
        # 计算 Torch 模型与 onnx 模型输出的绝对误差
        abs_error = np.abs(torch_output - onnx_output).mean()
        print(f'Mean Error: {abs_error:.2f}')
        return onnx_model

对于 ResNet50 而言，Torch 模型的推断耗时为 172.67 ms，onnx 模型的推断耗时为 36.56 ms，onnx 模型的推断耗时仅为 Torch 模型的 21.17%

到此这篇关于PyTorch 模型 onnx 文件导出及调用详情的文章就介绍到这了,更多相关PyTorch文件导出内容请搜索脚本之家以前的文章或继续浏览下面的相关文章希望大家以后多多支持脚本之家！

您可能感兴趣的文章:

Python爬虫实现HTTP网络请求多种实现方式
这篇文章主要介绍了Python爬虫实现HTTP网络请求多种实现方式,文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值,需要的朋友可以参考下
2020-06-06
详解python中list的使用
这篇文章主要介绍了python中list的使用，文中通过示例代码介绍的非常详细，对大家的学习或者工作具有一定的参考学习价值，需要的朋友们下面随着小编来一起学习学习吧
2019-03-03
python 列表,数组,矩阵两两转换tolist()的实例
下面小编就为大家分享一篇python 列表,数组,矩阵两两转换tolist()的实例，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2018-04-04
Python的Tkinter点击按钮触发事件的例子
今天小编就为大家分享一篇Python的Tkinter点击按钮触发事件的例子，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2019-07-07
PYQT5 实现给listwidget的滚动条添加滚动信号
这篇文章主要介绍了PYQT5 实现给listwidget的滚动条添加滚动信号，具有很好的参考价值，希望对大家有所帮助。一起跟随小编过来看看吧
2021-03-03
在Python中比较列表中的相邻元素的几种方法
在Python中,我们可以通过多种方式来对比列表中的相邻项,我们没有看到任何直接或间接的应用程序来比较相邻的元素,例如确定最近的趋势,优化用户体验,股票市场分析等等,本文将探讨在Python中如何比较列表中的相邻元素的几种方法,需要的朋友可以参考下
2025-01-01
Python通过Django实现用户注册和邮箱验证功能代码
这篇文章主要介绍了Python通过Django实现用户注册和邮箱验证功能代码，具有一定借鉴价值，需要的朋友可以参考下。
2017-12-12
Python利用pyHook实现监听用户鼠标与键盘事件
这篇文章主要介绍了Python利用pyHook实现监听用户鼠标与键盘事件,很有实用价值的一个技巧,需要的朋友可以参考下
2014-08-08
Python实现接口自动化测试的方法详解
Python接口自动化测试是一种高效、可重复的软件质量验证方法,尤其在现代软件开发中,它已经成为不可或缺的一部分,本文将深入探讨如何使用Python进行接口自动化测试,文中通过代码示例介绍的非常详细,需要的朋友可以参考下
2024-08-08
自学python用什么系统好
在本篇文章里小编给大家整理了一篇关于学python用什么系统好的相关文章，有兴趣的朋友们可以学习下。
2020-06-06

PyTorch 模型 onnx 文件导出及调用详情

目录

前言

基本用法

高级 API

相关文章

最新评论

大家感兴趣的内容

最近更新的内容

常用在线小工具