【PYTORCH ＞＞ ONNX ＞＞ TENSORRT】自定义plug-in算子解析onnx模型生成引擎详细步骤

博客概要

网络上可以找到许多如何在TensorRT引擎中自己添加plug-in算子的博客，但是他们的共同特点都是仅仅记录了一小部分细节问题，导致读者无法全面的了解这一套流程的详细方案，甚至有一些博客可能由于TensorRT的版本问题，甚至调用了一些不支持的API，导致我在进行实践的时候走了不少弯路，在这里本博客从代码层面进行对整个流程和一些细节关键点进行了详细全面的总结，对自己是个备忘的同时也希望能帮助到阅读本博客的童鞋。

为什么要写plug-in

一般而言是因为tensorrt原生不支持一些算子，为了整个网络的其他算子仍然可以被转化为tensorrt引擎所以要进行用户自定义，对于更高阶的用户，也许想把tensorrt支持的算子替换掉，这仍是可行的，但后者不需要进行本博客一样完全的模型转换过程而仅仅是在tensorrt项目内部进行核函数的撰写和算子的手动替换(因为tensorrt算子和onnx模型节点已经建立好了内部的替换映射关系)。

pytorch TensorRT环境版本

pytorch: 1.9.0 + cu111 .
TensorRT: cuda11.4 + cudnn8.2.1.32 + tensorrt 8.4.1.5 .
安装好上述环境后，对于自定义plug-in还是不充足的，还需要在github上下载tensorrt的项目源代码,链接如下：
链接: TensorRT项目链接
由于该项目添加了一系列如protobuf等多个子模块项目链接，手动点击下载是相对麻烦的，因此需要参照该项目的readme进行git clone操作，具体命令如下：

git clone -b master https://github.com/nvidia/TensorRT TensorRT
cd TensorRT
git submodule update --init --recursive

第一行命令和第三行命令都是需要下载的，如果太慢作者的解决方案是给linux挂了一个VPN代理，其他方式都是可以的这里就不再赘述了就希望大家可以各显神通啦。下载好了该项目后，参照readme进行编译即可，蛋疼的是他会继续在内部进行protobuf依赖项目下载，所以必须必须要有VPN！！！否则完全可能出现下载卡死的情况。这里作者还发现了一个点，在NVIDIA官方下载的tensorrt deb安装包安装出的libnvinfer_plugin.so动态库为版本为8.4.1，但是在github tensorrt项目中编译出的版本为8.2.0。8.4.1库内已经自带并注册好的plugin算子约63个，而8.2.0版本仅有40多个，因为不影响后续的开发工作具体缺少了哪一些算子没有去深究了，需要知道的是这样进行下去自己写的plugin是添加在8.2.0的动态库版本上面的，因为形成的新的动态库会直接替换掉旧的动态库。

TensorRT项目介绍

从github下载下来的TensorRT项目，可以大致分为最重要的三个部分：他们分别是

1. plugin 内部有一些tensorrt自带的开源plugin，有了这些文件，我们可以非常高效的依样画葫芦定义自己的plugin。
2. parser 如果用户准备在tensorrt内部使用API一层一层的把网络搭建起来，则该文件夹对于上述用户而言是没用的，但当你的模型是从onnx或者caffe等模型转换过来的，tensorrt并不知道该算子应该如何关联转换，是会在build engine阶段报错的，本项目由于是torch onnx trt转换，因此是需要去修改parser内部文件内容，使得trt认识我们自定义的算子在onnx计算图内对应的到底是什么东西，在这里也被称作op_type。
3. samples 学习使用，对于本博客可以忽略，但是这些代码是学习trt非常有用的开源工具，在trt项目最外层的cmakelist内，可以将对samples的编译关闭来加快编译时间
   
   
   完成上述准备工作后，我们要从源头pytorch开始出发了。

从torch出发自定义一个算子

本文章用gelu激活函数举例，在torch.nn中是有gelu的实现的，但当你导出定义的gelu模型为onnx，该模型会让你感觉好乱好乱，可以认为onnx当前暂时是不支持gelu的(op version 11)，但是它很聪明，他知道gelu无非就是加减乘除一顿输出就完事了，因此onnx会给你一个这样的计算图：

实际上，pytorch内部我是这么去定义这个模型的：

class MyNet(nn.Module):
	def __init__(self):
		super(MyNet, self).__init__()
		self.gelu = nn.GELU()

	def forward(self, x):
		return self.gelu(self.gelu(x))

我们知道这样的onnx模型肯定是可以转换为trt模型的，因为该计算图内都是一些最基本的计算单元的组合。由于trt是一个闭源软件，作者不是特别肯定其内部会不会对这样的算子进行图融合操作(当然可以用nsight去看看trt核函数的调用堆栈去尽可能的推断一下，这与本文无关不再赘述)。如果trt内部做了而且做的还蛮好，ok那本文到此结束 XD。假设trt执行的就是上述计算图，因为gelu并不是一个计算密度很高的算子，因此该算子带来的latency主要是由访存体现的，对于global memory一读一写是耗时的最大贡献来源，所以我们希望该算子可以长成这样：

在pytorch中自定义算子并且被onnx正确识别

直接上pytorch官方自定义算子链接：extending pytorch
由于我们仅仅是需要有一个特定的算子符号，而不是需要将该模型拿来训练，因此我们是只需要对于forward和symbolic函数进行重写，而backword函数是无需重写的。

import torch
import torch.nn as nn
from torch.autograd import Function

class MyGelu(Function):
	@staticmethod
	def