极智AI | 初识 TensorRT Plugin

极智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 初识​​​​2 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: // 重写部分成员函数 }; ...} //

写个 Makefile 去编译出 ​​.so​​：

CUDA_PATH = /usr/local/cuda/NVCC = $(CUDA_PATH)/bin/nvccTRT_PATH = /usr/lib/x86_64-linux-gnu/SM = 61CCFLAG = -std=c++11 -UNDEBUG -DDEBUG -O3 -arch=sm_$(SM)SOFLAG = $(CCFLAG) -sharedINCLUDE = -I. -I$(CUDA_PATH)/include -I$(TRT_PATH)/includeLDFLAG = -L$(CUDA_PATH)/lib64 -lcudart -L$(TRT_PATH)/lib -lnvinferSOURCE_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: testtest: make clean make python $(SOURCE_PY).PHONY: cleanclean: 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》​​

版权声明：本文内容由网络用户投稿，版权归原作者所有，本站不拥有其著作权，亦不承担相应法律责任。如果您发现本站中有涉嫌抄袭或描述失实的内容，请联系我们jiasou666@gmail.com 处理，核实后本网站将在24小时内删除侵权内容。