三. TensorRT基础入门-导出并分析ONNX

前言

自动驾驶之心推出的《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第三章—TensorRT 基础入门，一起来学习如何导出并分析 ONNX

课程大纲可以看下面的思维导图

在这里插入图片描述

0. 简述

本小节目标：学习 Pytorch 导出 ONNX 的方法以及 onnx-simplifier 的使用

这节我们学习第三章节第四小节—导出 ONNX、分析 ONNX，这里我们拆分成了几个小节来分析 ONNX，包括 ONNX 内部的 protobuf 怎么用以及它的数据表达形式是什么样的，导出 ONNX 时遇到不兼容的算子怎么处理等等内容

这个小节还是比较简单的，我们主要来学习下 Pytorch 模型怎么导出 ONNX，以及 onnx-simplifier 模块的使用

1. generate-onnx

源代码获取地址：https://github.com/kalfazed/tensorrt_starter

这个部分对应的案例主要是 3.1-generate-onnx，如下所示：

在这里插入图片描述

我们直接进入 src 代码文件中看下，先看下 example.py 文件，代码如下所示：

import torch
import torch.nn as nn
import torch.onnxclass Model(torch.nn.Module):def __init__(self, in_features, out_features, weights, bias=False):super().__init__()self.linear = nn.Linear(in_features, out_features, bias)with torch.no_grad():self.linear.weight.copy_(weights)def forward(self, x):x = self.linear(x)return xdef infer():in_features = torch.tensor([1, 2, 3, 4], dtype=torch.float32)weights = torch.tensor([[1, 2, 3, 4],[2, 3, 4, 5],[3, 4, 5, 6]],dtype=torch.float32)model = Model(4, 3, weights)x = model(in_features)print("result is: ", x)def export_onnx():input   = torch.zeros(1, 1, 1, 4)weights = torch.tensor([[1, 2, 3, 4],[2, 3, 4, 5],[3, 4, 5, 6]],dtype=torch.float32)model   = Model(4, 3, weights)model.eval() #添加eval防止权重继续更新# pytorch导出onnx的方式，参数有很多，也可以支持动态size# 我们先做一些最基本的导出，从netron学习一下导出的onnx都有那些东西torch.onnx.export(model         = model, args          = (input,),f             = "../models/example.onnx",input_names   = ["input0"],output_names  = ["output0"],opset_version = 12)print("Finished onnx export")if __name__ == "__main__":infer()export_onnx()

上述代码定义了一个简单的 PyTorch 模型，进行了推理和模型导出为 ONNX 的操作，我们可以使用 Netron 工具来查看下我们导出的 ONNX，如下图所示：

在这里插入图片描述

它其实是一个非常简单的 ONNX 架构，一个 input，一个 output，中间有一个 MatMul 的节点，右侧栏中可以看到各个节点的相关信息，比如 input 的 name 叫 input0，它的 type 是 float32，维度是 1x1x1x4

在代码中我们使用的是 torch.onnx.export 函数来将我们定义的 pytorch 模型导出为 onnx 模型 torch.onnx.export 函数是 PyTorch 用来将模型导出到 ONNX (Open Neural Network Exchange) 格式的工具。这个函数非常重要，因为它允许模型在不同的机器学习框架间进行转换和部署。这里博主询问了 ChatGPT 并记录了该函数中各个参数的介绍方便后续查看，下面是这个函数及其主要参数的详细介绍：（from ChatGPT）

def export(model: Union[torch.nn.Module, torch.jit.ScriptModule, torch.jit.ScriptFunction],args: Union[Tuple[Any, ...], torch.Tensor],f: Union[str, io.BytesIO],export_params: bool = True,verbose: bool = False,training: _C_onnx.TrainingMode = _C_onnx.TrainingMode.EVAL,input_names: Optional[Sequence[str]] = None,output_names: Optional[Sequence[str]] = None,operator_export_type: _C_onnx.OperatorExportTypes = _C_onnx.OperatorExportTypes.ONNX,opset_version: Optional[int] = None,do_constant_folding: bool = True,dynamic_axes: Optional[Union[Mapping[str, Mapping[int, str]], Mapping[str, Sequence[int]]]] = None,keep_initializers_as_inputs: Optional[bool] = None,custom_opsets: Optional[Mapping[str, int]] = None,export_modules_as_functions: Union[bool, Collection[Type[torch.nn.Module]]] = False,autograd_inlining: Optional[bool] = True,
) -> None:

model (torch.nn.Module): 这是你要导出的 PyTorch 模型。它应该是一个已经训练好的模型。（重要）
args (tuple 或 torch.Tensor): 这是传递给模型的输入数据。这个参数非常关键，因为它用于推导模型中张量的形状。如果你的模型接受多个输入，可以通过元组传递每个输入。（重要）
f (str 或 io.BytesIO): 指定导出的 ONNX 文件的保存路径或一个二进制文件。如果是字符串，它是文件的路径；如果是 BytesIO 对象，模型将被保存到这个内存中的对象。
export_params (bool, 默认为 True): 指定是否导出模型的权重。设为 True 会一同导出模型参数，否则只导出模型结构。
verbose (bool, 默认为 False): 如果设为 True，在导出过程中会打印出模型的详细信息，这有助于调试。
training (torch.onnx.TrainingMode 或 bool): 用于指定模型是处于训练模式 (TrainingMode.TRAINING) 还是评估模式 (TrainingMode.EVAL)。这会影响某些层（如批归一化层）的导出方式。
input_names (list[str]): 为输入张量指定名字，这些名字在 ONNX 图中会用作节点的标识。
output_names (list[str]): 为输出张量指定名字，这有助于在 ONNX 模型中识别输出节点。
operator_export_type (torch.onnx.OperatorExportTypes): 控制某些 PyTorch 操作如何转换为 ONNX 操作。例如，ONNX_ATEN_FALLBACK 使得如果某个操作无法被标准 ONNX 操作覆盖，可以使用 PyTorch 自定义操作。
opset_version (int): 指定导出的 ONNX 模型使用的操作集版本。不同的版本可能支持不同的操作。默认是使用 PyTorch 支持的最新的 ONNX 版本。（重要）
do_constant_folding (bool, 默认值为 True): 此参数控制是否应用常量折叠优化。常量折叠是一种优化技术，它会预先计算那些所有输入都是常量的操作，并将这些操作替换为预计算的常量节点。
dynamic_axes (dict): 允许指定哪些维度可以是动态的。这在处理变长的序列或者批量大小可变的情况时非常有用。（重要）
keep_initializers_as_inputs (Optional[bool]): 这个参数决定 initializers（如权重和偏差）是否也应该作为输入列在 ONNX 计算图中。将此设置为 True 可以增加与某些期望权重作为图输入的 ONNX 推理引擎的旧版本的兼容性。如果 opset_version < 9，那么 initializers 必须作为图的输入。
custom_opsets (Optional[Mapping[str, int]]): 允许你为特定操作集定义版本号。这在使用非标准或自定义 ONNX 操作时特别有用，可以通过这个参数指定这些操作的版本。
export_modules_as_functions (Union[bool, Collection[Type[torch.nn.Module]]]): 这个参数可以设置为 True 或者一个包含模块类型的集合，用于导出时将这些模块作为 ONNX 中的函数进行管理。
autograd_inlining (Optional[bool]): 这个参数用于控制是否在导出过程中将自动微分操作内联到图中。如果设置为 True（默认值），则相关的自动微分计算会被内联进图中，这有可能会使导出的模型更简洁，但在某些情况下可能不希望这种行为。

下面是该函数的一个使用示例：

import torch
import torch.nn as nn# 定义模型
class SimpleModel(nn.Module):def __init__(self):super(SimpleModel, self).__init__()self.linear = nn.Linear(10, 5)def forward(self, x):return self.linear(x)# 实例化模型并设置为评估模式
model = SimpleModel().eval()# 创建一个示例输入
example_input = torch.randn(1, 10)# 导出模型
torch.onnx.export(model,               # 要导出的模型example_input,       # 模型的示例输入"model.onnx",        # 保存模型的文件名export_params=True,  # 导出模型的参数opset_version=11,    # 指定 ONNX opset 版本do_constant_folding=True,  # 是否执行常数折叠优化input_names=['input'],      # 输入的名字output_names=['output'],    # 输出的名字dynamic_axes={'input' : {0 : 'batch_size'},  # 指定批量大小是动态的'output' : {0 : 'batch_size'}})

我们再来看下 example_two_head.py，代码如下所示：

import torch
import torch.nn as nn
import torch.onnxclass Model(torch.nn.Module):def __init__(self, in_features, out_features, weights1, weights2, bias=False):super().__init__()self.linear1 = nn.Linear(in_features, out_features, bias)self.linear2 = nn.Linear(in_features, out_features, bias)with torch.no_grad():self.linear1.weight.copy_(weights1)self.linear2.weight.copy_(weights2)def forward(self, x):x1 = self.linear1(x)x2 = self.linear2(x)return x1, x2def infer():in_features = torch.tensor([1, 2, 3, 4], dtype=torch.float32)weights1 = torch.tensor([[1, 2, 3, 4],[2, 3, 4, 5],[3, 4, 5, 6]],dtype=torch.float32)weights2 = torch.tensor([[2, 3, 4, 5],[3, 4, 5, 6],[4, 5, 6, 7]],dtype=torch.float32)model = Model(4, 3, weights1, weights2)x1, x2 = model(in_features)print("result is: \n")print(x1)print(x2)def export_onnx():input    = torch.zeros(1, 1, 1, 4)weights1 = torch.tensor([[1, 2, 3, 4],[2, 3, 4, 5],[3, 4, 5, 6]],dtype=torch.float32)weights2 = torch.tensor([[2, 3, 4, 5],[3, 4, 5, 6],[4, 5, 6, 7]],dtype=torch.float32)model   = Model(4, 3, weights1, weights2)model.eval() #添加eval防止权重继续更新# pytorch导出onnx的方式，参数有很多，也可以支持动态sizetorch.onnx.export(model         = model, args          = (input,),f             = "../models/example_two_head.onnx",input_names   = ["input0"],output_names  = ["output0", "output1"],opset_version = 12)print("Finished onnx export")if __name__ == "__main__":infer()export_onnx()

这个示例代码与之前的代码相比，主要的区别在于模型中增加了一个额外的线性层和一个额外的输出，导出的 ONNX 如下图所示：

在这里插入图片描述

可以看到导出的 ONNX 有两个 head，与我们代码中设置的一样

我们继续看下一个 example_dynamic_shape.py，代码如下所示：

import torch
import torch.nn as nn
import torch.onnxclass Model(torch.nn.Module):def __init__(self, in_features, out_features, weights, bias=False):super().__init__()self.linear = nn.Linear(in_features, out_features, bias)with torch.no_grad():self.linear.weight.copy_(weights)def forward(self, x):x = self.linear(x)return xdef infer():in_features = torch.tensor([1, 2, 3, 4], dtype=torch.float32)weights = torch.tensor([[1, 2, 3, 4],[2, 3, 4, 5],[3, 4, 5, 6]],dtype=torch.float32)model = Model(4, 3, weights)x = model(in_features)print("result of {1, 1, 1 ,4} is ", x.data)def export_onnx():input   = torch.zeros(1, 1, 1, 4)weights = torch.tensor([[1, 2, 3, 4],[2, 3, 4, 5],[3, 4, 5, 6]],dtype=torch.float32)model   = Model(4, 3, weights)model.eval() #添加eval防止权重继续更新# pytorch导出onnx的方式，参数有很多，也可以支持动态sizetorch.onnx.export(model         = model, args          = (input,),f             = "../models/example_dynamic_shape.onnx",input_names   = ["input0"],output_names  = ["output0"],dynamic_axes  = {'input0':  {0: 'batch'},'output0': {0: 'batch'}},opset_version = 12)print("Finished onnx export")if __name__ == "__main__":infer()export_onnx()

上述代码和前面的示例相比最显著的变化在于 ONNX 导出部分，这里引入了动态维度的支持，导出的 ONNX 如下图所示：

在这里插入图片描述

我们可以对比之前导出的 ONNX，在没有动态维度支持时导出的 ONNX 中 batch 维度为固定数值 1，如果是动态 shape 如上图所示可以看到 batch 维度不再是一个固定的数值，而是动态可变化的。值得注意的是，我们在导出 ONNX 时一般只会让 batch 维度动态，宽高不动态

2. export-onnx

源代码获取地址：https://github.com/kalfazed/tensorrt_starter

这个部分对应的案例主要是 3.2-export-onnx，如下所示：

在这里插入图片描述

我们直接进入 src 代码文件中看下，先看下 sample_cbr.py 文件，代码如下所示：

import torch
import torch.nn as nn
import torch.onnxclass Model(torch.nn.Module):def __init__(self):super().__init__()self.conv1 = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3)self.bn1   = nn.BatchNorm2d(num_features=16)self.act1  = nn.ReLU()def forward(self, x):x = self.conv1(x)x = self.bn1(x)x = self.act1(x)return xdef export_norm_onnx():input   = torch.rand(1, 3, 5, 5)model   = Model()model.eval()# 通过这个案例，我们一起学习一下onnx导出的时候，其实有一些节点已经被融合了# 思考一下为什么batchNorm不见了file    = "../models/sample-cbr.onnx"torch.onnx.export(model         = model, args          = (input,),f             = file,input_names   = ["input0"],output_names  = ["output0"],opset_version = 15)print("Finished normal onnx export")if __name__ == "__main__":export_norm_onnx()

上面代码展示了一个简单的 PyTorch 模型，其中包括一个卷积层（Conv2d）、一个批量归一化层（BatchNorm2d）和一个激活函数（ReLU），以及如何将这个模型导出为ONNX格式。

导出的 ONNX 如下图所示：

在这里插入图片描述

那大家可以对上面导出的 ONNX 有所疑问，为什么 BatchNorm 层在导出的ONNX文件中不见了呢？那这其实因为当 PyTorch 模型被导出到 ONNX 格式时，可以发生所谓的图优化，其中一些连续的操作可能被融合为一个操作，以优化执行效率。尤其是当 Conv2d 和 BatchNorm2d 层相连时，这两个操作经常会被融合为一个单一的卷积操作。这种优化被称为卷积-批量归一化融合：(from ChatGPT)

原因: 在评估模式下，批量归一化的参数（均值、方差）是固定的，而卷积层的权重也是固定的。这意味着这两个层的操作可以组合成一个具有新权重的单一卷积操作，而无需改变输出结果。这样做可以减少运行时的计算负担，因为减少了运行的层的数量
效果: 这样的融合能有效减少模型的复杂性和提高执行效率，特别是在硬件加速环境下（如使用 GPU 或专用 AI 加速器）

总的来说，通过此案例我们需要了解到 ONNX 导出过程中的节点融合是一个常见的优化手段，用于提高模型的运行效率。这种优化是自动进行的，基于当前的模型结构和 opset 版本的支持。

我们来看下一个案例 sample_reshape.py 文件，代码如下所示：

import torch
import torch.nn as nn
import torch.onnx
import onnxsim
import onnxclass Model(torch.nn.Module):def __init__(self):super().__init__()self.conv1   = nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1)self.bn1     = nn.BatchNorm2d(num_features=16)self.act1    = nn.ReLU()self.conv2   = nn.Conv2d(in_channels=16, out_channels=64, kernel_size=5, padding=2)self.bn2     = nn.BatchNorm2d(num_features=64)self.act2    = nn.ReLU()self.avgpool = nn.AdaptiveAvgPool1d(1)self.head    = nn.Linear(in_features=64, out_features=10)def forward(self, x):x = self.conv1(x)x = self.bn1(x)x = self.act1(x)x = self.conv2(x)x = self.bn2(x)x = self.act2(x) x = torch.flatten(x, 2, 3)  # B, C, H, W -> B, C, L (这一个过程产生了shape->slice->concat->reshape这一系列计算节点, 思考为什么)# b, c, w, h = x.shape# x = x.reshape(b, c, w * h)# x = x.view(b, c, -1)x = self.avgpool(x)         # B, C, L    -> B, C, 1x = torch.flatten(x, 1)     # B, C, 1    -> B, Cx = self.head(x)            # B, L       -> B, 10return xdef export_norm_onnx():input   = torch.rand(1, 3, 64, 64)model   = Model()file    = "../models/sample-reshape.onnx"torch.onnx.export(model         = model, args          = (input,),f             = file,input_names   = ["input0"],output_names  = ["output0"],opset_version = 15)print("Finished normal onnx export")model_onnx = onnx.load(file)# 检查导入的onnx modelonnx.checker.check_model(model_onnx)# 使用onnx-simplifier来进行onnx的简化。# 可以试试把这个简化给注释掉，看看flatten操作在简化前后的区别# onnx中其实会有一些constant value，以及不需要计算图跟踪的节点# 大家可以一起从netron中看看这些节点都在干什么# print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")# model_onnx, check = onnxsim.simplify(model_onnx)# assert check, "assert check failed"onnx.save(model_onnx, file)if __name__ == "__main__":export_norm_onnx()

上述代码实现了一个包含多个层的神经网络，并展示了如何将 PyTorch 模型导出为 ONNX 格式，以及如何检查和简化 ONNX 模型，导出的 ONNX 模型如下图所示：

在这里插入图片描述

那看上面的 ONNX 模型大家可能会想为什么 torch.flatten 这个操作产生了 shape->slice->concat->reshape 这一系列计算节点呢？

这是因为在 PyTorch 模型中使用 torch.flatten 操作时，如果涉及到多个维度的平展，这个操作在转换到 ONNX 格式时通常会被分解成几个更基础的操作（如 Shape、Slice、Concat 和 Reshape），主要是因为 ONNX 需要明确地追踪数据的形状变化，让 ONNX 标准化其操作，以确保与不同的后端和硬件兼容性。下面是每个步骤作用的详细解释：(from ChatGPT)

Shape：这个操作获取张量的维度。在 ONNX 中，许多操作需要知道输入张量的具体形状来执行后续的操作，尤其是在维度转换或变化时。
Slice：这个操作用于从 Shape 得到的形状张量中提取特定的维度。在示例中 torch.flatten(x, 2, 3) 意味着要将第二维到第三维的尺寸合并。Slice 操作帮助提取这些维度的尺寸。
Concat：这个操作将 Slice 操作得到的尺寸值与其他维度的尺寸进行合并。在 flatten 操作中，合并的目的是创建一个新的尺寸数组，用于下一步的 Reshape。
Reshape：根据 Concat 得到的新尺寸数组，重新塑形张量。这实际上是实现 flatten 的最终步骤，将指定的多个维度合并为一个维度。

这个过程看起来比直接在 PyTorch 中使用 flatten 更复杂，因为 ONNX 需要更明确地表达维度变化，以确保模型的兼容性和可移植性。每一个步骤都是必要的，以在不同的平台和框架之间确保操作的一致性。

如果在将模型导出为 ONNX 时遇到性能问题或者想要优化这一过程，可以考虑是否有方法简化网络结构或者预处理步骤，或者使用 ONNX 的优化工具来尝试合并这些操作。

下面我们就通过 onnx-simplifier 来优化这个 ONNX 模型，优化后的 ONNX 模型如下图所示：

在这里插入图片描述

可以看到优化后的 ONNX 模型非常的干净，之前 flatten 的一系列操作直接变成了 reshape 一个节点搞定

补充：onnx-simplifier 是一个用于简化 ONNX 模型的工具。它通过合并冗余的操作、消除不必要的中间节点等方法，优化模型的计算图。onnx-simplifier 通常以命令行工具的形式提供，可以直接通过 Python 包管理器安装，使用简单的命令就可以对 ONNX 模型进行简化。例如：

pip install onnx-simplifier
python -m onnxsim input_model.onnx output_model.onnx

关于 onnx-simplifier 的更多细节可以查看 https://github.com/daquexian/onnx-simplifier

OK，我们来看最后一个案例程序 load_torchvision.py，代码如下所示：

import torch
import torchvision
import onnxsim
import onnx
import argparsedef get_model(type, dir):if type == "resnet":model = torchvision.models.resnet50()file  = dir + "resnet50.onnx"elif type == "vgg":model = torchvision.models.vgg11()file  = dir + "vgg11.onnx"elif type == "mobilenet":model = torchvision.models.mobilenet_v3_small()file  = dir + "mobilenetV3.onnx"elif type == "efficientnet":model = torchvision.models.efficientnet_b0()file  = dir + "efficientnetb0.onnx"elif type == "efficientnetv2":model = torchvision.models.efficientnet_v2_s()file  = dir + "efficientnetV2.onnx"elif type == "regnet":model = torchvision.models.regnet_x_1_6gf()file  = dir + "regnet1.6gf.onnx"return model, filedef export_norm_onnx(model, file, input):model.cuda()torch.onnx.export(model         = model, args          = (input,),f             = file,input_names   = ["input0"],output_names  = ["output0"],opset_version = 15)print("Finished normal onnx export")model_onnx = onnx.load(file)# 检查导入的onnx modelonnx.checker.check_model(model_onnx)# 使用onnx-simplifier来进行onnx的简化。print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")model_onnx, check = onnxsim.simplify(model_onnx)assert check, "assert check failed"onnx.save(model_onnx, file)def main(args):type        = args.typedir         = args.dirinput       = torch.rand(1, 3, 224, 224, device='cuda')model, file = get_model(type, dir)export_norm_onnx(model, file, input)if __name__ == "__main__":parser = argparse.ArgumentParser()parser.add_argument("-t", "--type", type=str, default="resnet")parser.add_argument("-d", "--dir", type=str, default="../models/")opt = parser.parse_args()main(opt)

上述代码的主要功能是导出不同类型的预训练神经网络模型到 ONNX 格式，并对它们进行简化。它使用了 PyTorch 的 torchvision 库来获取各种常用的神经网络模型，然后将它们转换为 ONNX 格式，并使用 onnx-simplifier 进行优化。代码整合了命令行参数支持，以便用户可以选择导出的模型类型和保存位置。

执行后导出的 ONNX 如下图所示：

在这里插入图片描述

大家可以导出其它的一些模型，利用 Netron 看看它们的网络结构都长什么样的，有什么不同

3. 补充-ONNX

以下内容均 copy 自：Jetson嵌入式系列模型部署-1

3.1 概念

onnx 可以理解为一种通用货币，开发者可以把自己开发训练好的模型保存为onnx文件，而部署工程师可以借助部署框架（如 tensorRT、openvino、ncnn 等）部署在不同的硬件平台上，而不必关系开发者使用的是哪一种框架
onnx 的本质是一种 protobuf 格式文件
protobuf 通过编译 onnx-ml.proto 文件得到 onnx-ml.pb.h 和 onnx-ml.pb.cc 用于 C++ 调用或 onnx_ml_pb2.py 用于 python 调用，如下图所示。如果本地 python 环境下安装了 onnx 第三方库，则在该库下可以找到 onnx_ml_pb2.py 文件

在这里插入图片描述

通过编译得到的 onnx-ml.pb.cc 和代码就可以操作 onnx 模型文件，实现对应的增删改
onnx-ml.proto 用于描述 onnx 文件是如何组成的，具有什么结构，它是 onnx 经常参照的东西，如下是 onnx-ml.proto 部分内容，参考自：https://github.com/shouxieai/tensorRT_Pro/blob/main/onnx/onnx-ml.proto

在这里插入图片描述

3.2 组成

onnx 文件组成如下图所示

在这里插入图片描述

model：表示整个 onnx 模型，包括图结构和解析器版本、opset 版本、导出程序类型
- opset 版本即operator 版本号即 pytorch 的 op (操作算子)版本
model.graph：表示图结构，通常是 Netron 可视化工具中看到的结构
model.graph.node：表示图结构中所有节点如 conv、bn、relu 等
model.graph.initializer：权重数据大都存储在这里
model.graph.input：模型的输入
model.graph.output：模型的输出

总结

本次课程我们主要学习了利用 torch.onnx.export 函数将 Pytorch 模型导出为 ONNX 模型，并学习了利用 onnx-simplifier 来优化我们导出的 ONNX 模型。

OK，以上就是第 4 小节有关 ONNX 导出的全部内容了，下节我们来学习剖析 ONNX 架构并理解 Protobuf，敬请期待😄