5.4.tensorRT基础(2)-学习第一个插件的编写
目录
-
- 前言
- 1. 插件
- 2. 补充知识
- 总结
前言
杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。
本次课程学习 tensorRT 基础-学习第一个插件的编写
课程大纲可看下面的思维导图
1. 插件
这节我们学习插件的实现
插件实现-重点：
1. 如何在 pytorch 里面导出一个插件
2. 插件解析时如何对应，在 onnx parser 中如何处理
3. 插件的 creator 实现
4. 插件的具体实现，继承自 IPluginV2DynamicExt
5. 插件的序列化和反序列化
首先我们提供了一个 gen-onnx.py 文件，用于在 pytorch 中导出一个插件，内容如下：
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import torch.autograd
import os
class MYSELUImpl(torch.autograd.Function):
# reference: https://pytorch.org/docs/1.10/onnx.html#torch-autograd-functions
@staticmethod
def symbolic(g, x, p):
print("==================================call symbolic")
return g.op("MYSELU", x, p,
g.op("Constant", value_t=torch.tensor([3, 2, 1], dtype=torch.float32)),
attr1_s="这是字符串属性",
attr2_i=[1, 2, 3],
attr3_f=222
)
@staticmethod
def forward(ctx, x, p):
return x * 1 / (1 + torch.exp(-x))
class MYSELU(nn.Module):
def __init__(self, n):
super().__init__()
self.param = nn.parameter.Parameter(torch.arange(n).float())
def forward(self, x):
return MYSELUImpl.apply(x, self.param)
class Model(nn.Module):
def __init__(self):
super().__init__()
self.conv = nn.Conv2d(1, 1, 3, padding=1)
self.myselu = MYSELU(3)
self.conv.weight.data.fill_(1)
self.conv.bias.data.fill_(0)
def forward(self, x):
x = self.conv(x)
x = self.myselu(x)
return x
# 这个包对应opset11的导出代码，如果想修改导出的细节，可以在这里修改代码
# import torch.onnx.symbolic_opset11
print("对应opset文件夹代码在这里：", os.path.dirname(torch.onnx.__file__))
model = Model().eval()
input = torch.tensor([
# batch 0
[
[1, 1, 1],
[1, 1, 1],
[1, 1, 1],
],
# batch 1
[
[-1, 1, 1],
[1, 0, 1],
[1, 1, -1]
]
], dtype=torch.float32).view(2, 1, 3, 3)
output = model(input)
print(f"inference output = \n{output}")
dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(
model,
# 这里的args，是指输入给model的参数，需要传递tuple，因此用括号
(dummy,),
# 储存的文件路径
"workspace/demo.onnx",
# 打印详细信息
verbose=True,
# 为输入和输出节点指定名称，方便后面查看或者操作
input_names=["image"],
output_names=["output"],
# 这里的opset，指，各类算子以何种方式导出，对应于symbolic_opset11
opset_version=11,
# 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1
# 通常，我们只设置batch为动态，其他的避免动态
dynamic_axes={
"image": {0: "batch", 2: "height", 3: "width"},
"output": {0: "batch", 2: "height", 3: "width"},
},
# 对于插件，需要禁用onnx检查
enable_onnx_checker=False
)
print("Done.!")
其中 Model 类还是和之前一样的 conv+activation，不过这次激活函数替换成了我们自己实现的一个 MYSELU，其 forward 是调用了 MYSELUImpl 这个类，它是继承自 torch.autograd.function，只要是插件你就要继承自它，在这个插件实现类中你需要实现一个静态 forward 方法，它接收 ctx，x，param 三个参数。
你想实现自定义算子的话，你可以自己实现 forward 甚至是 backward，另外 forwad 中的部分不会被跟踪，也就是意味着不会生成节点，那真正生成节点靠的是 symbolic 这个静态函数，具体可参考 pytorch 官方文档的 https://pytorch.org/docs/1.10/onnx.html#torch-autograd-functions
它接收 g，x，param 三个参数，通过 g.op 创建一个节点，我们来导出一下 onnx，出现了如下错误：
提示 onnx.export() 没有 enable_onnx_checker 这个参数，应该是我的 torch 版本的原因，换了我的笔记本电脑没问题
导出完成之后还打印了推理结果，待会和 C++ 的推理结果做对比，导出的 onnx 如下图所示：
可以看到激活函数是我们自定义实现的 MYSELU，它还有 Constant 这个节点输出，是一些常量值，接下来我们执行下 C++ 的代码来解析下这个 onnx 文件，看能不能拿到正确的结果，执行下 make run
可以看到正常解析并执行完推理了，同时推理结果和 pytorch 保持一致，说明我们的插件能够被正常执行，main.cpp 中除了包含自己编译的 onnx-tensorrt-release-8.0 之外没有任何不同，那 MYSELU 这个算子是如何被解析的呢？
在 builtin_op_importers.cpp 最后我们注册了一个叫做 MYSELU 的节点，内容如下：
DEFINE_BUILTIN_OP_IMPORTER(MYSELU)
{
printf("\033[31m=======================call MYSELU==============\033[0m\n");
OnnxAttrs attrs(node, ctx);
const std::string pluginName{node.op_type()};
const std::string pluginVersion{attrs.get<std::string>("plugin_version", "1")};
const std::string pluginNamespace{attrs.get<std::string>("plugin_namespace", "")};
LOG_INFO("Searching for plugin: " << pluginName << ", plugin_version: " << pluginVersion << ", plugin_namespace: " << pluginNamespace);
nvinfer1::IPluginCreator* creator = importPluginCreator(pluginName, pluginVersion, pluginNamespace);
ASSERT(creator && "Plugin not found, are the plugin name, version, and namespace correct?", ErrorCode::kUNSUPPORTED_NODE);
const nvinfer1::PluginFieldCollection* fieldNames = creator->getFieldNames();
// Field data needs to be type erased, we use fieldData for temporary allocations.
string_map<std::vector<uint8_t>> fieldData{};
std::vector<nvinfer1::PluginField> fields = loadFields(fieldData, attrs, fieldNames, ctx);
const auto plugin = createPlugin(getNodeName(node), creator, fields);
ASSERT(plugin && "Could not create plugin", ErrorCode::kUNSUPPORTED_NODE);
std::vector<nvinfer1::ITensor*> pluginInputs{};
for (auto& input : inputs)
{
pluginInputs.emplace_back(&convertToTensor(input, ctx));
}
LOG_INFO("Successfully created plugin: " << pluginName);
auto* layer = ctx->network()->addPluginV2(pluginInputs.data(), pluginInputs.size(), *plugin);
ctx->registerLayer(layer, getNodeName(node));
RETURN_ALL_OUTPUTS(layer);
}
正是这个节点的注册才能让 onnx 解析器正常解析这个节点，那这个节点解析出来了怎么对应到执行的呢，我们通过 importPluginCreator() 将 pluginName 即 node.op_type() 即 MYSELU 当参数传入创建一个 creator 得到一个插件的创建器，通过 creator 去读取 onnx 对应的属性，然后将对应的属性丢到 createPlugin() 这个函数去创建我们的 plugin，然后通过 ctx->network->addPluginV2() 得到对应的 layer，最后通过 registerLayer 完成我们插件的注册。
那这个只是插件的创建和注册，并没有说明它是怎么和我们的执行代码关联起来的，我们来看下 myselu-plugin.hpp 和 myselu-plugin.cpp 这两个文件，首先头文件它要实现 PluginCreater，要实现获取插件的名字、版本号、创建插件实例、反序列化插件…，内容如下：
#ifndef CUSTOM_MYSELU_PLUGIN_H
#define CUSTOM_MYSELU_PLUGIN_H
#include "NvInferPlugin.h"
#include 创建插件实例将参数序列化为文件，不再走 onnx，因此反序列化你也需要实现。其次头文件中还有插件类的实现，这是我们关注的重点，首先它继承自 IPluginV2DynamicExt，我们要实现两种构造方式，通过参数来构造以及通过数据来构造，然后你还需要实现一些特定的函数，比如获取输出的数量、输出的数据类型、输出的维度…，其中 enqueue 才是我们执行时使用的函数，它不仅仅在运行时才执行，在编译阶段 tensorRT 会以实际方式在设备上运行以此选择哪个配置是最合适的。
对应的实现在 CPP 文件中，其内容如下：
#include "myselu-plugin.hpp"
#include "NvInfer.h"
#include 通过 supportsFormatCombination 函数来告诉插件支持的数据格式和类型，enqueue 具体的 inference 就是在调用核函数，具体实现如下：
#include 到此为止，我们大体走了一遍插件实现的流程，好难呀
2. 补充知识
关于插件你需要知道：(from 杜老师)
知识点
1. 插件主要是继承自 IPluginV2DynamicExt 后实现特定接口即可
实现插件要点
1. 导出 onnx 的时候，需要为 module 增加 symbolic 函数
- 参照这里：https://pytorch.org/docs/1.13/onnx.html#torch-autograd-functions
- g.op 对应的名称，需要与下面解析器的名称对应
2. src/onnx-tensorrt-release-8.0/builtin_op_importers.cpp:5094行，添加对插件 op 的解析
- DEFINE_BUILTIN_OP_IMPORTER(MYSELU)
- 注意解析时采用的名称要匹配上 src/myselu-plugin.cpp:15行
3. src/myselu-plugin.cpp:183行，创建 MySELUPluginCreator，插件创建器
- 实际注册时，注册的是创建器，交给 tensorRT 管理
- REGISTER_TENSORRT_PLUGIN(MySELUPluginCreator);
- src/myselu-plugin.cpp:23行
4. src/myselu-plugin.cpp:42行，定义插件类 MySELUPlugin
- Creator 创建器来实例化 MySELUPlugin 类
5. 正常使用该 onnx 即可
实现插件的两个阶段
1. 编译阶段
- 通过 MySELUPluginCreator::createPlugin 创建 plugin
- 期间会调用 MySELUPlugin::clone 克隆插件
- 调用 MySELUPlugin::supportsFormatCombination 判断该插件所支持的数据格式和类型
在这里我们告诉引擎，本插件可以支持什么类型的推理
可以支持多种，例如 fp32、fp16、int8 等等
- 调用 MySELUPlugin::getOutputDimension 获取该层的输出维度是多少
- 调用 MySELUPlugin::enqueue 进行性能测试（不是一定会执行）
如果支持多种，则会在多种里面进行实际测试，选择一个性能最好的配置
- 调用 MySELUPlugin::configurePlugin 配置插件格式
告诉你目前这个层所采用的数据格式和类型
- 调用 MySELUPlugin::serialize 将该层的参数序列化储存为 trtmodel 文件
2. 推理阶段
- 通过 MySELUPluginCreator::deserializePlugin 反序列化插件参数进行创建
- 期间会调用 MySELUPlugin::clone 克隆插件
- 调用 MySELUPlugin::configurePlugin 配置当前插件使用的数据类型和格式
- 调用 MySELUPlugin::enqueue 进行推理
总结
本节课程学习了插件的编写，首先你要在 pytorch 中能够正确导出 onnx，自定义算子需要增加 symbolic 函数，然后你需要去 onnx-tensorrt 源码中的 buildin_op_importers 中添加对插件 op 的解析，接下来你要实现两个类，一个是插件创建类 Creator，另外一个是插件具体实现类，继承自 IPluginV2DynamicExt，其中具体推理代码还需要编写 CUDA 核函数完成，最后就是通过宏定义来注册插件。
感觉写插件的难度还是大，目前还没有完整的走过一遍流程
而且一般来说一些简单的算子实现 tensorRT 应该都帮你搞定了，只有一些复杂的算子才需要写插件，流程虽然说是这么个流程，但是感觉复杂算子的插件难度上升得有点高，简单的都还没玩明白
非常硬核的技术活