极智AI | 初识 TensorRT Plugin

  欢迎关注我的公众号 [极智视界]，获取我的更多笔记分享

  大家好，我是极智视界，本文介绍一下 初识 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，希望我的分享能对你的学习有一点帮助。

 【极智视界传送】

《极智AI | 初识 TensorRT Plugin》