TensorRT Plugin实战

追求源于热爱688

已于 2023-11-10 21:44:50 修改

阅读量1.8k

点赞数 2

分类专栏： TensorRT总结文章标签：计算机视觉神经网络 YOLO 深度学习 cnn

于 2023-02-17 18:03:08 首次发布

本文链接：https://blog.csdn.net/sinat_41886501/article/details/129091397

版权

TensorRT总结专栏收录该内容

2 篇文章 0 订阅

订阅专栏

TensorRT-Plugin小试牛刀

注：需要在linux下编译，win10没有编译成功

代码已上传至github，github链接：https://github.com/ljx6666/TensorRT-Plugin

整体流程：

先安装TRT库；
修改onnx，加入plugin；
编译对应TRT版本的源代码(win10上会遇到很多坑，linux会相对顺利一点)；
修改TRT源代码，加入对应plugin的cpp、h和cu文件，然后还要注册，其他的cpp和h文件需要改一下；
编译改完的TRT源代码，生成的.so和.a文件放到TRT库的lib目录中，替换掉原来的lib目录；
使用trtexec工具，将改过的onnx生成engine；
验证阶段，在调用自定义的plugin时，需要在反序列化engine(deserializeCudaEngine函数)前，加入头文件#include “NvInferPlugin.h”和bool didInitPlugins = initLibNvInferPlugins(&gLogger, “ ”); ，并在Cmakelists.txt中加入1行target_link_libraries(yolox libnvinfer_plugin.so)。

第一步：
·需要安装TensorRT库，选择8.2.0.6版本的，其他版本的也可以尝试下。

第二步：
在onnx中修改网络的激活函数，为了增加TRT中不支持的op类型，使用下面的脚本将LeakyRelu的op_type改为Custom：
在这里插入图片描述

最后尝试使用trtexec工具将修改后放入onnx模型导出，报错。。。
在这里插入图片描述

第三步：
·下载与安装的TensorRT版本一致的源代码进行编译
最好是安装一样的版本，比如TensorRT库是8.2.0.6的，TensorRT源码也要是8.2.0.6的。不是同一版本的没试过，所以不知道好不好使。
下载命令：

git clone -b release/8.2.0.6 https://github.com/nvidia/TensorRT TensorRT  // 8.2.0.6可以替换成别的版本，如7.1
或：
git clone -b master https://github.com/nvidia/TensorRT TensorRT               // 默认就是8.2.0.6版本的

cd TensorRT
git submodule update --init --recursive                           // 下载trt源码中的第三方库如onnx等，有时候网不好，需要多下几回，才能下全。

如下图所示，输入git submodule update --init --recursive 命令后没有输出了，证明下全了。
在这里插入图片描述

在make之前，需要在Cmakelists.txt里修改1行代码，如下图所示，需要将TRT_LIB_DIR路径指定为TensorRT库的lib目录：
在这里插入图片描述

编译命令：

mkdir build
cd build
cmake .. 
make -j10

正常情况下是很顺利，不会报错的，编译完之后会在build目录下生成.so和.a文件。做这步的目的是为了验证此TensorRT的源码是否可以编译成功，因为加入新的plugin之后，还得编译一遍TensorRT的源码，在此之前，先把TensorRT源码编译的坑先填平了。。。
如果编译TRT其他版本源码，如7.1版本，注意7.1版本要求CUDA版本是11.0或10.2，如果要用其他CUDA版本编译，需要在cmake那步更改命令为：

cmake .. -DCUDA_VERSION=11.2    // 这里使用的CUDA版本是11.2，可替换成你自己想要的版本

除此之外，TRT7.1版本的源码使用的第三方库默认的protobuf版本是3.0.0，如果在编译的过程中下载不下来，可提前利用科学上网下载下来tar.gz，此时需要修改对应的CMakeLists.txt，如下图所示：
在这里插入图片描述

第四步：
打开TRT源代码，在plugin目录中照着LeakyReluPlugin复制1份，将h文件和cpp文件里面的lRelu统统替换成Custom，把文件名也改了，如下图所示。

图1

图2 lCustomPlugin.cpp(除了图中展示的部分，其他地方还有需要改成Custom的)

为了后续方便注册，自定义插件Custom类需要继承nvinfer1::IPluginV2DynamicExt接口，另外需要增加一些重写方法和属性，如图3、4、5：
在这里插入图片描述

图3 lCustomPlugin.h前部分

在这里插入图片描述

图4 lCustomPlugin.h中间部分

在这里插入图片描述

图5 lCustomPlugin.h后部分

在plugin/common/kernel.h中，加入图6中的那一行。
在这里插入图片描述

图6 kernel.h

在plugin/common/kernels中，需要写一个lCustom.cu文件，仿照lRelu.cu写，在文件里面把lRelu替换成Custom，如图7所示，最终的Custom plugin函数实现在pCustomKernel()里，可以在此处修改成自己想要的plugin函数代码。
在这里插入图片描述

图7 lCustom.cu

还得需要注册plugin，在InferPlugin.cpp中加入图8、9两行代码：
在这里插入图片描述

图8 InferPlugin.cpp

在这里插入图片描述

图9 InferPlugin.cpp

以及在plugin/CMakeLists.txt中添加如图10这一行：
在这里插入图片描述

图10 plugin/CMakeLists.txt

最后需要实现onnx结点和TRT插件的映射关系，改动如下：
在这里插入图片描述

图11 builtin_op_importers.cpp

第五步：
编译改完的TRT源码，编译的命令与第三步一致，编译完之后将build目录下的.so和.a文件拷贝到TRT库中的lib目录下，选择“全部替换”。

第六步：
利用trtexec工具，将改完的onnx生成engine，结果如下：
在这里插入图片描述

图12 custom.engine

第七步：
如何调用自定义plugin呢？
在这里插入图片描述

图13 yolox_s.onnx修改完节点后模型对比

如果在推理代码中什么都不改，即使engine生成成功，在调用带有自定义插件的engine时，也会出错，出错截图如图14。
在这里插入图片描述

图14 推理代码未改动时，编译成功了，但在运行时还是出错了，报错如上

在网上百度了解决方案，见参考链接2，需要在推理代码中在对应位置加入图15、16，其中图16需要在deserializeCudaEngine函数前加上。除此之外，cmakelists.txt也需要加入1行，如图17所示。
在这里插入图片描述

图15 在yolox.cpp中加入NvInferPlugin.h头文件

在这里插入图片描述

图16 在deserializeCudaEngine函数前加上初始化plugin函数的调用

在这里插入图片描述

图17 CMakeLists.txt中需要链接到libnvinfer_plugin.so

最后，就可以愉快的调用自定义plugin模型推理啦！hhh
在这里插入图片描述

图18 用带有自定义plugin的yolox_s.engine推理结果

以上是介绍了一个TensorRT-Plugin的小示例，在实战中我将yolox的前处理resize等节点改成Plugin，方法同上。由于resize节点本身就存在于onnx转TensorRT的算子中，并不需要再编写Plugin的C++代码了，只需要在onnx模型中加入Shape、Gather、Concat和Resize节点。
图19为将yolox的前处理放进模型中的可视化结果，在原来模型的基础上增加了Shape、Gather、Concat和Resize节点，其中Concat节点和Resize节点可视化如图20所示，Concat节点的输出即为Resize节点输出的形状。加入Shape、Gather和Concat节点的目的是为了获取到网络输入的动态形状。至于为什么不在Resize节点的sizes参数里直接指定形状为[1, 3, 640, 640]，是因为直接指定会报错。
在这里插入图片描述

图19 左为yolox_s原模型，右为加入前处理节点后的yolox_s模型

在这里插入图片描述

图20 左为Concat节点，右为Resize节点

其实将resize的C++代码转成TensorRT-Plugin的目的是为了减少每帧的处理时间，以达到提升性能的目的。不过我测试完时间之后，发现使用resize-plugin反而让每帧的处理时间变长了。。。猜测原因是输入的图像尺寸变大了（因为这时输入的是原图），在resize之前还需要将nhwc转成nchw。输入图像尺寸变大，导致转换维度这一函数耗时也变长。除此之外，使用nsight systems查看resize节点的推理时长也不短。。。于是改变一下优化策略，改成将yolox的前处理和后处理函数全部写成CUDA C内核函数的方式去减少耗时，提升性能。CUDA内核函数的编写在以后的博客中会陆续介绍。

参考链接：
1、https://zhuanlan.zhihu.com/p/492144628
2、https://blog.csdn.net/a2824256/article/details/121262135?spm=1001.2101.3001.6650.11&utm_medium=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7ERate-11-121262135-blog-102723545.pc_relevant_default&depth_1-utm_source=distribute.pc_relevant.none-task-blog-2%7Edefault%7ECTRLIST%7ERate-11-121262135-blog-102723545.pc_relevant_default&utm_relevant_index=14