极智AI | 初识 TensorRT Plugin

Posted 极智视界

tags:

篇首语:本文由小常识网(cha138.com)小编为大家整理,主要介绍了极智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》


以上是关于极智AI | 初识 TensorRT Plugin的主要内容,如果未能解决你的问题,请参考以下文章

极智AI | 讲解 TensorRT 显式batch 和 隐式batch

极智AI | 讲解 TensorRT Constant 算子

极智AI | 讲解 TensorRT 怎么实现 torch.select 层

我的NVIDIA开发者之旅 - 极智AI | TensorRT 中 Layer 和 Tensor 的区别

极智AI | Attention 中 torch.chunk 的 TensorRT 实现

我的NVIDIA开发者之旅 - 极智AI | TensorRT API 构建模型推理流程