5.5.tensorRT基础(2)-封装插件过程，并实现更容易的插件开发

前言

杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。

本次课程学习 tensorRT 基础-封装插件过程，并实现更容易的插件开发

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

1. 插件封装

这节课我们来学习集成插件，也就是简化插件的实现，因为之前的插件实现起来太过复杂了

相比于之前插件的实现多了两个 onnxplugin.cpp 和 onnxplugin.hpp 文件，这是杜老师实现的关于插件的具体封装，我们只需要实现 easy-plugin.cu 这个文件，整个过程十分简单，我们先 make run 执行一下：

图1-1 make run出错

博主尝试 debug 调试了一下，先定位到 builtin_op_importers.cpp 文件的 4661 行，即

LOG_INFO("Successfully created plugin: " << pluginName);

// 4661行
auto* layer = ctx->network()->addPluginV2(pluginInputs.data(), pluginInputs.size(), *plugin);

ctx->registerLayer(layer, getNodeName(node));
RETURN_ALL_OUTPUTS(layer);

然后再定位到 onnxplugin.cpp 的 307 行，当返回输出的数量时直接崩了

int TRTPlugin::getNbOutputs() const noexcept{
    return config_->num_output_;
}

调试输出：ERROR: Couldn't get registers: No such process.

博主经过一番折腾，终于发现了是导出的 onnx 存在问题，使用杜老师提供的 onnx 没有问题，使用自己导出的就存在问题，最终发现博主直接使用的上节课程导出的 onnx，这是不行的，因为封装后的插件的 g.op_type() 必须是 Plugin 才行，因此要用本节课程的 gen-onnx.py 产生 onnx，否则就会导致找不到对应的 registers 而失败，执行成功后如下所示：

图1-2 make run执行成功

可以发现执行的效果和我们上节课没有封装的效果一样，我们重点关注下杜老师是怎么对插件进行封装的吧

我们回过头来看下代码，easy-plugin.cu 的具体内容如下：

#include "onnx-tensorrt/onnxplugin.hpp"

using namespace ONNXPlugin;

static __device__ float sigmoid(float x){
    return 1 / (1 + expf(-x));
}

static __global__ void MYSELU_kernel_fp32(const float* x, float* output, int edge) {

    int position = threadIdx.x + blockDim.x * blockIdx.x;
	if(position >= edge) return;

    output[position] = x[position] * sigmoid(x[position]);
}

class MYSELU : public TRTPlugin {
public:
	SetupPlugin(MYSELU);

	virtual void config_finish() override{
		printf("\033[33minit MYSELU config: %s\033[0m\n", config_->info_.c_str());
		printf("weights count is %d\n", config_->weights_.size());
	}

	int enqueue(const std::vector<GTensor>& inputs, std::vector<GTensor>& outputs, const std::vector<GTensor>& weights, void* workspace, cudaStream_t stream) override{
		
		int n = inputs[0].count();
		const int nthreads = 512;
		int block_size = n < nthreads ? n : nthreads;
		int grid_size = (n + block_size - 1) / block_size;

		MYSELU_kernel_fp32 <<<grid_size, block_size, 0, stream>>> (inputs[0].ptr<float>(), outputs[0].ptr<float>(), n);
		return 0;
	}
};

RegisterPlugin(MYSELU);

可以发现它就实现了两个函数，一个是 config_finish() 用于打印（非必需），一个是 enqueue() 用于具体推理执行，插件的实现中很多变量都相同，因此可以给定默认值，封装起来。

那接下来我们简单了解下它具体是怎么封装的，首先是 MYSELU 继承自 TRTPlugin，通过 SetupPlugin(MYSELU) 初始化插件，RegisterPlugin(MYSELU) 来注册插件。杜老师在 builtin_op_importers.cpp 中实现了一个通用的 Plugin，那这个通用的 Plugin 可以通过解析 g.op_type() 为 Plugin 的所有插件，这也是为什么 gen-onnx.py 中的 g.op() 第一个参数必须是 Plugin，也解释了博主第一次执行报错的原因。

对于 attribute 属性我们直接用 JSON 文件来控制，解析的时候直接按照 JSON 格式解析就行，这比之前 attribute 解析要方便很多，同时我们也使用了 name_s 字符串属性来表明插件的名字 “MYSELU”，我们可以看下导出的 onnx 长什么样子

图1-3 demo.onnx

可以看到我们的插件节点就叫做 Plugin，它的名字叫做 MYSELU，它的属性信息全部存储在 info 这个 json 中，这个是我们导出的情况。

接下来我们看它是如何实现的，通过宏定义 RegisterPlugin(MYSELU) 来实现一个通用默认的插件，而 TRTPlugin 继承自 IPluginV2DynamicEtx，也实现了很多默认函数，比如输出类型和输入一致，默认输出数目为 1 等等，如果你的插件不是这样的，你可以直接覆盖它，非常方便，我们对 enqueue 也做了一个封装，我们对序列化和反序列化也做了一个默认实现，因此实际插件的使用只需要实现 enqueue 就能解决绝大部分问题。

更多具体细节内容需要大家自行查看相关代码实现了

2. 补充知识

关于插件的封装，你需要知道：(from 杜老师)

知识点

1. 对插件进行了封装，使用起来更简单

2. 在 onnx-tensorrt 中添加了 onnxplugin.cpp，实现对 IPluginV2DynamicExt 的封装

3. 在 onnx-tensorrt/builtin_op_importers.cpp:5095 行，添加了 Plugin 的解析支持

• DEFINE_BUILTIN_OP_IMPORTER(Plugin)
• 使得只要名字是 Plugin 的节点，都可以解释到该函数上
• 在代码中，为通用插件提供了支持，使得使用者只需要继承简单的插件接口即可完成需求

4. 在 gen-onnx.py 导出时，symbolic 函数返回时，g.op 返回的永远都是 Plugin 这个名字，然后 name_s 指定为自己注册的插件名称，info_s 则传递为 json 字符串，那么复合属性就可以轻易得到支持

封装后的插件实现

1. 导出 onnx 时，按照 gen-onnx.py 在 symbolic 函数返回时，指定 g.op 的 name 为 Plugin

2. 指定 g.op 中 name_s 属性为注册的插件名称，对应后续插件类的类名

3. 指定 g.op 中 info_ 属性为需要读取的复合属性，字符串。通常可以传递 json，使得属性再复制都可以，避免使用官方的方式

4. 创建 easy-plugin.cu 文件，定义自己的类并继承自 ONNXPlugin::TRTPlugin

5. 实现需要的函数：

• config_finish[非必要]：配置完成函数
• 当插件配置完毕时调用，可以在其中拿到各种属性，例如 info、weights 等
• new_config[非必要]：实例化一个配置对象
• 可以自定义 LayerConfig 类并返回，也可以直接使用 LayerConfig 类
• 这个函数的最大作用，是配置本插件支持的数据格式和类型。比如 fp32 和 fp16 的支持等
• getOutputDimensions[非必要]，获取该插件输出的 shape 大小，默认取第一个输入的大小
• 对应于原始插件的 getOutputDimensions 函数
• enqueue[必要]，插件推理过程
• 插件的实际推理过程，该函数可能在编译和推理阶段数次调用

6. 注册插件，使用 RegisterPlugin 宏

• RegisterPlugin(MYSELU)
• 格式是 RegisterPlugin(类名)

7. 好了，可以使用插件了

总结

本节课程学习了插件的封装。首先杜老师在 builtin_op_importers.cpp 添加了对通用插件 Plugin 的解析支持，凡是叫做 Plugin 的节点，都能被正常解析，因此在导出 onnx 的时候需要指定 g.op_type() 为 Plugin，name_s 指定为自己注册的插件名称，info_s 则传递属性值为 json 字符串。

接下来在 onnxplugin.cpp 中实现对 IPluginV2DynamicExt 的封装，由于插件的很多属性都是通用的，因此实现了默认函数，比如序列化、反序列化等等，我们只需要实现 enqueue 然后调用封装的接口就可以完成整个插件的开发工作。

那经过这么一层封装，确实省事了不少呀，具体封装的细节就没去看了，先能写个插件用起来再说吧