大家好，我是极智视界，本文介绍一下 初识 TensorRT Plugin。
TensorRT 构建模型推理一般有三种方式：(1) 使用框架自带的 TensorRT 接口，如 TF-TRT、Torch-TRT；(2) 使用 Parser 前端解释器，如 TF / Torch / … -> ONNX -> TensorRT；(3) 使用 TensorRT 原生 API 搭建网络。如果你不做啥动作，这三种方式最后在构建 Engine 的时候，都会去调用 TensorRT 算子 API。那么你有没有想过一个问题：要是你想调用的算子，TensorRT 原生不支持，那该如何是好？这个时候，一般会催生出两种解决方法：(1) 将 TensorRT 原生算子 排列组合 后达到我们想要的样子；(2) 使用 TensorRT Plugin 进行算子构建。说到这里，就把 TensorRT Plugin 引出来了。
1 TensorRT Plugin 初识
从序言的介绍大致可以了解到 TensorRT Plugin 的作用是什么，实现原生不支持的算子 是 Plugin 最基础的能力，当然它还可以做更多事情，比如手动融合 TensorRT 没有自动融合的层 或 块。总结来说，TensorRT Plugin 的功能主要有以下几点：
- (1) 实现 TensorRT 原生不支持的算子 或 块；
- (2) 手动融合 TensorRT 没有自动融合的算子 或 块；
- (3) 替换你认为性能不够的 算子 或 块；
TensorRT Plugin 从实现形式来说，需要自己使用 CUDA C 来编写 CUDA Kernel，然后以 .so 的形式 嵌入到咱们网络的 TensorRT network 构建中，这可不是一件容易的事情。我认为其中的难度主要体现在：
- (1) 写出高性能的 CUDA C Kernel 本身并不是一件简单的事情；
- (2) 自己写的 Kernel，若没有经过系统测试，并没有原生算子那么可靠；
所以，编写 TensorRT Plugin 一定是一个有门槛的工作。且使用 TensorRT Plugin 还有一些其他限制：
- (1) Plugin 算子和 TensorRT 原生算子之间不能进行自动融合，要做融合的话，你只能在 Plugin 里把 数据映射 关系写进去；
- (2) 可能在 Plugin 节点前后新增 格式转换的算子，会增加一定时间开销；
既然 TensorRT Plugin 有这么多缺陷和限制，咱们为啥还要用呢？我觉得出发点主要有两个：
- (1) 没办法啊 ==> 原生算子搞不出来啊，不用 Plugin 我能怎么办，难道要我换算子去重新训练吗？
- (1) 爽啊 ==> 使用 Plugin 会更加的 彻底，操控感十足，原生 API 构建就像开自动挡的车，而 Plugin 构建就像开手动挡的车，你应该知道，追求极致体验的赛车 可都是手动挡的。
2 TensorRT Plugin 简单示例
咱们来列举个案例：给输入张量的所有元素加上同一个常量。
先写 AddScalarPlugin.cu：
// AddScalarPlugin.cu // 用于计算的 kernel __global__ void addScalarKernel(const float *input, float *output, const float scalar, const int nElement) { const int index = blockIdx.x * blockDim.x + threadIdx.x; if (index >= nElement) return; float _1 = input[index]; float _2 = _1 + scalar; output[index] = _2; } namespace nvinfer1{ ... class AddScalarPlugin: public IPluginV2DynamicExt{ privite: // 一些自定义的数据 public: // 重写部分成员函数 }; class AddScalarPluginCreator: public IPluginCreator{ privite: // 一些自定义的数据 public: // 重写部分成员函数 }; ... } // namespace nvinfer1
写个 Makefile 去编译出 .so：
CUDA_PATH = /usr/local/cuda/ NVCC = $(CUDA_PATH)/bin/nvcc TRT_PATH = /usr/lib/x86_64-linux-gnu/ SM = 61 CCFLAG = -std=c++11 -UNDEBUG -DDEBUG -O3 -arch=sm_$(SM) SOFLAG = $(CCFLAG) -shared INCLUDE = -I. -I$(CUDA_PATH)/include -I$(TRT_PATH)/include LDFLAG = -L$(CUDA_PATH)/lib64 -lcudart -L$(TRT_PATH)/lib -lnvinfer SOURCE_CU = $(shell find . -name '*.cu') SOURCE_PY = $(shell find . -name '*.py') OBJ = $(shell find . -name *.o) DEP = $(OBJ:.o=.d) -include $(DEP) all: $(SOURCE_CU:%.cu=%.so) %.so: %.o $(NVCC) $(SOFLAG) $(LDFLAG) -o $@ $^ %.o: %.cu $(NVCC) $(CCFLAG) $(INCLUDE) -M -MT $@ -o $(@:.o=.d) $< $(NVCC) $(CCFLAG) $(INCLUDE) -Xcompiler -fPIC -o $@ -c $< .PHONY: test test: make clean make python $(SOURCE_PY) .PHONY: clean clean: rm -rf ./*.d ./*.o ./*.so ./*.plan
有了 .so 后，来看怎么调用集成进 TensorRT network：
... trt.init_libnvinfer_plugins(logger, '') ctypes.cdll.LoadLibrary('AddScalar.so') for creator in trt.get_plugin_registry().plugin_creator_list: if creator.name == 'AddScalar': parameterList = [] parameterList.append(trt.PluginField("scalar", np.float32(scalar), trt.PluginFieldType.FLOAT32)) plugin = creator.create_plugin(creator.name, trt.PluginFieldCollection(parameterList)) pluginLayer = network.add_plugin_v2([inputT0], plugin) ...
这样就完成了从 cu编写 -> make编译 -> so集成 的完整过程。
好了，以上分享了 初始 TensorRT Plugin，希望我的分享能对你的学习有一点帮助。
