pytorch模型(.pt)转onnx模型(.onnx)的方法

推理模型部署(一)：ONNX runtime 实践 - 知乎 (zhihu.com) （推荐，有从零构建一个小onnx模型的示例）

ONNX简介-CSDN博客（推荐）

深度学习模型转换与部署那些事(含ONNX格式详细分析) (bindog.github.io) （推荐）

Pytorch导出ONNX及模型可视化教程_onnx模型可视化-CSDN博客（推荐）

Python PyTorch export用法及代码示例 - 纯净天空 (vimsky.com) （推荐，torch.onnx.export函数的各个参数讲解的很清楚）

pytorch模型(.pt)转onnx模型(.onnx)的方法详解(1)_pt转onnx-CSDN博客

深度学习模型转换与部署那些事(含ONNX格式详细分析)_为什么要转onnx模型-CSDN博客

onnx.proto：onnx/onnx/onnx.proto at main · onnx/onnx · GitHub

onnx中每个计算节点的属性、输入和输出可以参考：onnx/docs/Operators.md at main · onnx/onnx · GitHub

onnx模型可视化以及pytorch算子与onnx节点对应关系_onnx gather算子-CSDN博客

ONNX是什么

（摘自ONNX简介）

ONNX: Open Neural Network Exchange.

在某一任务中将Pytorch或者TensorFlow模型转化为ONNX模型(ONNX模型一般用于中间部署阶段)，然后再拿转化后的ONNX模型进而转化为我们使用不同框架部署需要的类型。典型的几个线路：

Pytorch -> ONNX -> TensorRT
Pytorch -> ONNX -> TVM
TF – onnx – ncnn

ONNX相当于一个翻译的作用，这也是为什么ONNX叫做Open Neural Network Exchange。

protobuf序列化数据结构协议

（摘自ONNX简介）

利用Pytorch我们可以将model.pt转化为model.onnx格式的权重，在这里onnx充当一个后缀名称，model.onnx就代表ONNX格式的权重文件，这个权重文件不仅包含了权重值，也包含了神经网络的网络流动信息以及每一层网络的输入输出信息和一些其他的辅助信息。

ONNX既然是一个文件格式，那么我们就需要一些规则去读取它，或者写入它，ONNX采用的是protobuf这个序列化数据结构协议去存储神经网络权重信息。Protobuf是一种平台无关、语言无关、可扩展且轻便高效的序列化数据结构的协议，可以用于网络通信和数据存储。

ONNX结构分析

（摘自深度学习模型转换与部署那些事）

对于ONNX的了解，很多人可能仅仅停留在它是一个开源的深度学习模型标准，能够用于模型转换及部署但是对于其内部是如何定义这个标准，如何实现和组织的，却并不十分了解，所以在转换模型到ONNX的过程中，对于出现的不兼容不支持的问题有些茫然。

ONNX结构的定义基本都在这一个onnx.proto文件里面了，如何你对protobuf不太熟悉的话，可以先简单了解一下再回来看这个文件。当然我们也不必把这个文件每一行都看明白，只需要了解其大概组成即可，有一些部分几乎不会使用到可以忽略。

这里我把需要重点了解的对象列出来

ModelProto
GraphProto
NodeProto
AttributeProto
ValueInfoProto
TensorProto

我用尽可能简短的语言描述清楚上述几个Proto之间的关系：当我们将ONNX模型load进来之后，得到的是一个ModelProto，它包含了一些版本信息，生产者信息和一个非常重要的GraphProto；在GraphProto中包含了四个关键的repeated数组，分别是node(NodeProto类型)，input(ValueInfoProto类型)，output(ValueInfoProto类型)和initializer(TensorProto类型)，其中node中存放着模型中的所有计算节点，input中存放着模型所有的输入节点，output存放着模型所有的输出节点，initializer存放着模型所有的权重；那么节点与节点之间的拓扑是如何定义的呢？非常简单，每个计算节点都同样会有input和output这样的两个数组(不过都是普通的string类型)，通过input和output的指向关系，我们就能够利用上述信息快速构建出一个深度学习模型的拓扑图。最后每个计算节点当中还包含了一个AttributeProto数组，用于描述该节点的属性，例如Conv层的属性包含group，pads和strides等等，具体每个计算节点的属性、输入和输出可以参考这个Operators.md文档。

需要注意的是，刚才我们所说的GraphProto中的input输入数组不仅仅包含我们一般理解中的图片输入的那个节点，还包含了模型当中所有权重。举个例子，Conv层中的W权重实体是保存在initializer当中的，那么相应的会有一个同名的输入在input当中，其背后的逻辑应该是把权重也看作是模型的输入，并通过initializer中的权重实体来对这个输入做初始化(也就是把值填充进来)

修改ONNX模型

（摘自深度学习模型转换与部署那些事）

解决问题的最好办法是从根源入手，也就是从算法同学那边的模型代码入手，我们需要告诉他们问题出在哪里，如何修改。但是也有一些情况是无法通过修改模型代码解决的，或者与其浪费那个时间，不如我们部署工程师直接在ONNX模型上动刀解决问题。

还有一种更dirty的工作是，我们需要debug原模型和转换后的ONNX模型输出结果是否一致(误差小于某个阈值)，如果不一致问题出现在哪一层，现有的深度学习框架我们有很多办法能够输出中间层的结果用于对比，而据我所知，ONNX中并没有提供这样的功能；这就导致了我们的debug工作极为繁琐

所以如果有办法能够随心所欲的修改ONNX模型就好了。要做到这一点，就需要了解上文所介绍的ONNX结构知识了。

比如说我们要在网络中添加一个节点，那么就需要先创建相应的NodeProto，参照文档设定其的属性，指定该节点的输入与输出，如果该节点带有权重那还需要创建相应的ValueInfoProto和TensorProto分别放入graph中的input和initializer中，以上步骤缺一不可。

经过一段时间的摸索和熟悉，我写了一个小工具onnx-surgery并集成了一些常用的功能进去，实现的逻辑非常简单，也非常容易拓展。代码比较简陋，但是足以完成一些常见的修改操作。

PyTorch模型转ONNX模型

（摘自Pytorch导出ONNX及模型可视化教程）

torch.onnx.export函数实现了Pytorch模型导出到ONNX模型，在pytorch1.10.2中，torch.onnx.export函数参数如下：

def export(model, args, f, export_params=True, verbose=False, training=TrainingMode.EVAL,
        input_names=None, output_names=None, operator_export_type=None,
        opset_version=None, _retain_param_name=None, do_constant_folding=True,
        example_outputs=None, strip_doc_string=None, dynamic_axes=None,
        keep_initializers_as_inputs=None, custom_opsets=None, enable_onnx_checker=None,
        use_external_data_format=None)

大多数参数使用默认配置即可，下面对常用的几个参数进行介绍：

torch.onnx.export(
    model, # 需要转换的网络模型
    args, # ONNX模型输入，通常为 tuple 或 torch.Tensor
    f, # ONNX模型导出路径
    input_names=None, # 按顺序定义ONNX模型输入结点名称，格式为：list of str，若不指定，会使用默认名字
    output_names=None, # 按顺序定义ONNX模型输出结点名称，格式为：list of str，若不指定，会使用默认名字
    opset_version=11 # opset版本，地平线目前仅支持设置为 10 or 11
)

其它参数的介绍可参考官方torch.onnx.export()函数手册。

程序示例1

（摘自深度学习模型转换与部署那些事）

import torch
import torchvision

dummy_input = torch.randn(10, 3, 224, 224, device='cuda')
model = torchvision.models.alexnet(pretrained=False).cuda()

input_names = [ "actual_input_1" ] + [ "learned_%d" % i for i in range(16) ]
output_names = [ "output1" ]
file_path = "./alexnet.onnx"

torch.onnx.export(model, dummy_input, file_path, verbose=True, input_names=input_names, output_names=output_names)

程序示例2

（摘自Pytorch导出ONNX及模型可视化教程）

该节内容主要包括单输入网络构建、模型导出生成ONNX格式、导出的ONNX模型有效性验证三个部分。

import torch.nn as nn
import torch
import numpy as np
import onnx
import onnxruntime

# -----------------------------------#
#   定义一个简单的单输入网络   
# -----------------------------------#
class MyNet(nn.Module):
    def __init__(self, num_classes=10):
        super(MyNet, self).__init__()
        self.features = nn.Sequential(
            nn.Conv2d(3, 32, kernel_size=3, stride=1, padding=1),   # input[3, 28, 28]  output[32, 28, 28]          
            nn.BatchNorm2d(32),
            nn.ReLU(inplace=True),
            nn.Conv2d(32, 64, kernel_size=3, stride=2, padding=1),  # output[64, 14, 14]
            nn.BatchNorm2d(64),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=2)                             # output[64, 7, 7]
        )

        self.fc = nn.Linear(64 * 7 * 7, num_classes)

    def forward(self, x):
        x = self.features(x)
        x = torch.flatten(x, start_dim=1)
        x = self.fc(x)
        return x

# -----------------------------------#
#   导出ONNX模型函数
# -----------------------------------#
def model_convert_onnx(model, input_shape, output_path):
    dummy_input = torch.randn(1, 3, input_shape[0], input_shape[1])
    input_names = ["input1"]        # 导出的ONNX模型输入节点名称
    output_names = ["output1"]      # 导出的ONNX模型输出节点名称

    torch.onnx.export(
        model,
        dummy_input,
        output_path,
        verbose=False,          # 如果指定为True，在导出的ONNX中会有详细的导出过程信息description
        keep_initializers_as_inputs=False,  # 若为True，会出现需要warning消除的问题
        opset_version=11,       # 版本通常为10 or 11
        input_names=input_names,
        output_names=output_names,
    )


if __name__ == '__main__':
    model = MyNet()
    # print(model)
    # 建议将模型转成 eval 模式
    model.eval()
    # 网络模型的输入尺寸
    input_shape = (28, 28)      
    # ONNX模型输出路径
    output_path = './MyNet.onnx'

    # 导出为ONNX模型
    model_convert_onnx(model, input_shape, output_path)
    print("model convert onnx finsh.")

    # -----------------------------------#
    #   复杂模型可以使用下面的方法进行简化   
    # -----------------------------------#
    # import onnxsim
    # MyNet_sim = onnxsim.simplify(onnx.load(output_path))
    # onnx.save(MyNet_sim[0], "MyNet_sim.onnx")

    # -----------------------------------------------------------------------#
    #   第一轮ONNX模型有效性验证，用来检查模型是否满足 ONNX 标准   
    #   这一步是必要的，因为无论模型是否满足标准，ONNX 都允许使用 onnx.save 存储模型，
    #   我们都不会希望生成一个不满足标准的模型~
    # -----------------------------------------------------------------------#
    onnx_model = onnx.load(output_path)
    onnx.checker.check_model(onnx_model)
    print("onnx model check_1 finsh.")

    # ----------------------------------------------------------------#
    #   第二轮ONNX模型有效性验证，用来验证ONNX模型与Pytorch模型的推理一致性   
    # ----------------------------------------------------------------#
    # 随机初始化一个模型输入，注意输入分辨率
    x = torch.randn(size=(1, 3, input_shape[0], input_shape[1]))
    # torch模型推理
    with torch.no_grad():
        torch_out = model(x)
    print(torch_out)            # tensor([[-0.5728,  0.1695, ..., -0.3256,  1.1357, -0.4081]])
    # print(type(torch_out))      # <class 'torch.Tensor'>

    # 初始化ONNX模型
    ort_session = onnxruntime.InferenceSession(output_path)
    # ONNX模型输入初始化
    ort_inputs = {ort_session.get_inputs()[0].name: x.numpy()}
    # ONNX模型推理
    ort_outs = ort_session.run(None, ort_inputs)
    # print(ort_outs)             # [array([[-0.5727689 ,  0.16947027,  ..., -0.32555276,  1.13574252, -0.40812433]], dtype=float32)]
    # print(type(ort_outs))       # <class 'list'>，里面是个numpy矩阵
    # print(type(ort_outs[0]))    # <class 'numpy.ndarray'>
    ort_outs = ort_outs[0]        # 把内部numpy矩阵取出来，这一步很有必要

    # print(torch_out.numpy().shape)      # (1, 10)
    # print(ort_outs.shape)               # (1, 10)

    # ----------------------------------------------------------------#
    # 比较实际值与期望值的差异，通过继续往下执行，不通过引发AssertionError
    # 需要两个numpy输入
    # ----------------------------------------------------------------#
    np.testing.assert_allclose(torch_out.numpy(), ort_outs, rtol=1e-03, atol=1e-05)
    print("onnx model check_2 finsh.")

程序示例3

（摘自Pytorch导出ONNX及模型可视化教程）

该节内容主要包括多输入网络构建、模型导出生成ONNX格式、导出的ONNX模型有效性验证三个部分。

import torch.nn as nn
import torch
import numpy as np
import onnx
import onnxruntime

# -----------------------------------#
#   定义一个简单的双输入网络   
# -----------------------------------#
class MyNet_multi_input(nn.Module):
    def __init__(self, num_classes=10):
        super(MyNet_multi_input, self).__init__()
        self.conv1 = nn.Conv2d(3, 32, kernel_size=3, stride=2, padding=1)   # input[3, 28, 28]  output[32, 14, 14]
        self.bn1 = nn.BatchNorm2d(32)
        self.relu1 = nn.ReLU(inplace=True)

        self.conv2 = nn.Conv2d(1, 16, kernel_size=3, stride=2, padding=1)   # input[1, 28, 28]  output[32, 14, 14]
        self.bn2 = nn.BatchNorm2d(16)
        self.relu2 = nn.ReLU(inplace=True)

        self.fc = nn.Linear(48 * 14 * 14, num_classes)

    def forward(self, x, y):
        x = self.relu1(self.bn1(self.conv1(x)))
        y = self.relu2(self.bn2(self.conv2(y)))
        z = torch.cat((x, y), 1)
        z = torch.flatten(z, start_dim=1)
        z = self.fc(z)
        return z

# -----------------------------------#
#   导出ONNX模型函数
# -----------------------------------#
def multi_input_model_convert_onnx(model, input_shape, output_path):
    dummy_input1 = torch.randn(1, 3, input_shape[0], input_shape[1])
    dummy_input2 = torch.randn(1, 1, input_shape[0], input_shape[1])
    input_names = ["input1", "input2"]        # 导出的ONNX模型输入节点名称
    output_names = ["output1"]      # 导出的ONNX模型输出节点名称

    torch.onnx.export(
        model,
        (dummy_input1, dummy_input2),
        output_path,
        verbose=False,          # 如果指定为True，在导出的ONNX中会有详细的导出过程信息description
        keep_initializers_as_inputs=False,  # 若为True，会出现需要warning消除的问题
        opset_version=11,       # 版本通常为10 or 11
        input_names=input_names,
        output_names=output_names,
    )

if __name__ == '__main__':
    multi_input_model = MyNet_multi_input()
    # print(multi_input_model)
    # 建议将模型转成 eval 模式
    multi_input_model.eval()
    # 网络模型的输入尺寸
    input_shape = (28, 28)      
    # ONNX模型输出路径
    multi_input_model_output_path = './multi_input_model.onnx'

    # 导出为ONNX模型
    multi_input_model_convert_onnx(multi_input_model, input_shape, multi_input_model_output_path)
    print("multi_input_model convert onnx finsh.")

    # -----------------------------------#
    #   复杂模型可以使用下面的方法进行简化   
    # -----------------------------------#
    # import onnxsim
    # multi_input_model_sim = onnxsim.simplify(onnx.load(multi_input_model_output_path))
    # onnx.save(multi_input_model_sim[0], "multi_input_model_sim.onnx")

    # -----------------------------------------------------------------------#
    #   第一轮ONNX模型有效性验证，用来检查模型是否满足 ONNX 标准   
    #   这一步是必要的，因为无论模型是否满足标准，ONNX 都允许使用 onnx.save 存储模型，
    #   我们都不会希望生成一个不满足标准的模型~
    # -----------------------------------------------------------------------#
    onnx_model = onnx.load(multi_input_model_output_path)
    onnx.checker.check_model(multi_input_model_output_path)
    print("onnx model check_1 finsh.")

    # ----------------------------------------------------------------#
    #   第二轮ONNX模型有效性验证，用来验证ONNX模型与Pytorch模型的推理一致性   
    # ----------------------------------------------------------------#
    # 随机初始化一个模型输入，注意输入分辨率
    x = torch.randn(size=(1, 3, input_shape[0], input_shape[1]))
    y = torch.randn(size=(1, 1, input_shape[0], input_shape[1]))
    # torch模型推理
    with torch.no_grad():
        torch_out = multi_input_model(x, y)
    # print(torch_out)            # tensor([[-0.5728,  0.1695, ..., -0.3256,  1.1357, -0.4081]])
    # print(type(torch_out))      # <class 'torch.Tensor'>

    # 初始化ONNX模型
    ort_session = onnxruntime.InferenceSession(multi_input_model_output_path)
    # ONNX模型输入初始化
    ort_inputs = {ort_session.get_inputs()[0].name: x.numpy(), ort_session.get_inputs()[1].name: y.numpy()}
    # ONNX模型推理
    ort_outs = ort_session.run(None, ort_inputs)
    # print(ort_outs)             # [array([[-0.5727689 ,  0.16947027,  ..., -0.32555276,  1.13574252, -0.40812433]], dtype=float32)]
    # print(type(ort_outs))       # <class 'list'>，里面是个numpy矩阵
    # print(type(ort_outs[0]))    # <class 'numpy.ndarray'>
    ort_outs = ort_outs[0]        # 把内部numpy矩阵取出来，这一步很有必要

    # print(torch_out.numpy().shape)      # (1, 10)
    # print(ort_outs.shape)               # (1, 10)

    # ----------------------------------------------------------------#
    # 比较实际值与期望值的差异，通过继续往下执行，不通过引发AssertionError
    # 需要两个numpy输入
    # ----------------------------------------------------------------#
    np.testing.assert_allclose(torch_out.numpy(), ort_outs, rtol=1e-03, atol=1e-05)
    print("onnx model check_2 finsh.")

（十八）使用 DLA - 知乎 (zhihu.com)

posted @ 2024-07-05 09:04 Picassooo 阅读(6974) 评论(0) 收藏举报

刷新页面返回顶部

Picassooo

pytorch模型(.pt)转onnx模型(.onnx)的方法

ONNX是什么

protobuf序列化数据结构协议

ONNX结构分析

修改ONNX模型

PyTorch模型转ONNX模型

程序示例1

程序示例2

程序示例3

公告