什么是plugin & 有什么用?
TensorRT的一种机制,以.so的形式插入到网络中实现某些算子。
作用:
- 实现TensorRT不支持的层
- 替换性能不好的层
- 手动进行图优化算子融合
写Plugin就是自己写算子的CUDA kernel实现。
Plugin与其他layer之间无法自动进行算子融合,可能会在plugin前后加入reformating节点,增加开销。
建议先进行原生layer组合保证计算正确性,再尝试官方自带的Plugin是否满足要求,都不行再自己写plugin。
创建Plugin工作流程:
实现一个算子,对输入的张量每个元素加上一个常量
- 继承 IPluginV2DynamicExt 类实现一个Plugin 类
- 继承 IPluginCreator 类实现一个 PluginCreator 类
- 实现用于计算的 CUDA C++ kernel
- 将 Plugin 编译为 .so 保存
- 在 TenorRT 中加载和使用 Plugin
实现Plugin类
继承IPluginV2DynamicExt类
Plugin有V1和V2两个版本,V1已经弃用,V2分为:IPluginV2,IPluginV2Ext,IPluginV2IOExt,IPluginV2DynamicExt四种,第三种第四种最常用
class AddScalarPlugin : public IPluginV2DynamicExt // 定义AddScalarPlugin类,继承IPluginV2DynamicExt类
{
private:const std::string name_; //算子名称std::string namespace_; //算子所属的域struct{float scalar;} m_;public:AddScalarPlugin() = delete; //禁止默认构造函数AddScalarPlugin(const std::string &name, float scalar);AddScalarPlugin(const std::string &name, const void *buffer, size_t length); //构造函数~AddScalarPlugin();// Method inherited from IPluginV2const char *getPluginType() const noexcept override; //获取插件类型,noexcept表示该函数不会抛出异常,override表示该函数是虚函数const char *getPluginVersion() const noexcept override; //获取插件版本int32_t getNbOutputs() const noexcept override; //获取输出张量的数量int32_t initialize() noexcept override; //初始化插件void terminate() noexcept override; //终止插件,释放资源size_t getSerializationSize() const noexcept override; //获取序列化后的大小void serialize(void *buffer) const noexcept override; //序列化void destroy() noexcept override; //销毁插件,当context或engine被销毁时,插件也会被销毁void setPluginNamespace(const char *pluginNamespace) noexcept override; //设置插件的命名空间const char *getPluginNamespace() const noexcept override; //获取插件的命名空间//当我们的模型来自onnx的时候,命名空间,版本等信息会被保存在onnx模型中,这个函数就是用来获取这些信息的//一般不用我们自己设置,而是由onnx模型中的信息来设置//如果这些信息设置不对,会导致onnxparser解析模型的时候出错,无法识别插件// Method inherited from IPluginV2ExtDataType getOutputDataType(int32_t index, DataType const *inputTypes, int32_t nbInputs) const noexcept override;void attachToContext(cudnnContext *contextCudnn, cublasContext *contextCublas, IGpuAllocator *gpuAllocator) noexcept override;void detachFromContext() noexcept override;// Method inherited from IPluginV2DynamicExtIPluginV2DynamicExt *clone() const noexcept override;DimsExprs getOutputDimensions(int32_t outputIndex, const DimsExprs *inputs, int32_t nbInputs, IExprBuilder &exprBuilder) noexcept override;// getOutputDimensions,向TensorRT报告输出张量的形状,outputIndex是指输出张量的索引bool supportsFormatCombination(int32_t pos, const PluginTensorDesc *inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;// supportsFormatCombination,检查输入和输出张量的格式是否支持,pos是指输入张量的索引,inOut是指输入和输出张量的描述符, nbInputs是指输入张量的数量,nbOutputs是指输出张量的数量// 尽量多的支持格式组合,以便TensorRT可以选择最佳的格式组合void configurePlugin(const DynamicPluginTensorDesc *in, int32_t nbInputs, const DynamicPluginTensorDesc *out, int32_t nbOutputs) noexcept override;// configurePlugin,配置插件,in是指输入张量的描述符,nbInputs是指输入张量的数量,out是指输出张量的描述符,nbOutputs是指输出张量的数量// 在推理期前调用该函数,用于将插件中的动态维度转换为静态维度size_t getWorkspaceSize(const PluginTensorDesc *inputs, int32_t nbInputs, const PluginTensorDesc *outputs, int32_t nbOutputs) const noexcept override;// getWorkspaceSize,获取插件所需的工作空间大小,inputs是指输入张量的描述符,nbInputs是指输入张量的数量,outputs是指输出张量的描述符,nbOutputs是指输出张量的数量// 在推理期前调用该函数,用于计算插件所需的工作空间大小,向TensorRT报告工作空间的大小int32_t enqueue(const PluginTensorDesc *inputDesc, const PluginTensorDesc *outputDesc, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) noexcept override;// enqueue,执行插件的推理,inputDesc是指输入张量的描述符,outputDesc是指输出张量的描述符,inputs是指输入张量的数据,outputs是指输出张量的数据,workspace是指工作空间,stream是指CUDA流// 在推理期间调用该函数,用于执行插件的推理。不要在enqueue中调用cudaMalloc或cudaFree等CUDA API,会造成性能下降// 原因我猜是因为前面getworkspaceSize已经分配了空间,如果这里再进行分配,会使之前针对内存分配做的优化失效
protected:// To prevent compiler warnings,使用using声明,将基类的成员函数引入到子类中,避免编译器警告using nvinfer1::IPluginV2::enqueue;using nvinfer1::IPluginV2::getOutputDimensions;using nvinfer1::IPluginV2::getWorkspaceSize;using nvinfer1::IPluginV2Ext::configurePlugin;
};
实现PluginCreator类
继承IPluginCreator类
class AddScalarPluginCreator : public IPluginCreator
// 定义一个AddScalarPluginCreator类,继承于IPluginCreator,PluginCreator是一个工厂类,用于创建Plugin
{
private:static PluginFieldCollection fc_;static std::vector<PluginField> attr_;std::string namespace_;public:AddScalarPluginCreator();~AddScalarPluginCreator();const char *getPluginName() const noexcept override;const char *getPluginVersion() const noexcept override;const PluginFieldCollection *getFieldNames() noexcept override;IPluginV2DynamicExt *createPlugin(const char *name, const PluginFieldCollection *fc) noexcept override;// 接受一个插件名称和插件属性集合,返回一个新的插件实例IPluginV2DynamicExt *deserializePlugin(const char *name, const void *serialData, size_t serialLength) noexcept override;// 接受一个插件名称和序列化数据,返回一个新的插件实例void setPluginNamespace(const char *pluginNamespace) noexcept override;// 设置插件的命名空间const char *getPluginNamespace() const noexcept override;// 获取插件的命名空间
};
实现kernel函数
// kernel for GPU
__global__ void addScalarKernel(const float *input, float *output, const float scalar, const int nElement)
// cuda中global关键字修饰函数表示该函数必须由CPU调用,GPU运行
{const int index = blockIdx.x * blockDim.x + threadIdx.x;//cuda中kernel函数内置变量blockIdx表示目前执行该kernel的block信息,threadIdx表示执行该kernel的thread信息if (index >= nElement) // 如果越界就返回,否则会出现内存访问错误return; //cuda中kernel不允许返回值,但是return可以用来提前结束函数float _1 = input[index];float _2 = _1 + scalar;output[index] = _2;
}int32_t AddScalarPlugin::enqueue(const PluginTensorDesc *inputDesc, const PluginTensorDesc *outputDesc, const void *const *inputs, void *const *outputs, void *workspace, cudaStream_t stream) noexcept
{WHERE_AM_I();int nElement = 1;for (int i = 0; i < inputDesc[0].dims.nbDims; ++i){nElement *= inputDesc[0].dims.d[i];}dim3 grid(CEIL_DIVIDE(nElement, 256), 1, 1), block(256, 1, 1);addScalarKernel<<<grid, block, 0, stream>>>(reinterpret_cast<const float *>(inputs[0]), reinterpret_cast<float *>(outputs[0]), m_.scalar, nElement);return 0;
}
编译
include ../include/Makefile.incSOURCE_CU = $(shell find . -name '*.cu' 2>/dev/null)
SOURCE_PY = $(shell find . -name '*.py' 2>/dev/null)
OBJ = $(shell find . -name *.o 2>/dev/null)
DEP = $(OBJ:.o=.d)
TARGET_SO = $(SOURCE_CU:.cu=.so)-include $(DEP)all: $(TARGET_SO)%.so: %.o$(NVCC) $(SOFLAG) $(LDFLAG) -o $@ $+
# nvcc是指定编译器,-shared是指定生成动态链接库,-o是指定生成的动态链接库的名字,$+是指定生成动态链接库的目标文件%.o: %.cu$(NVCC) $(CUFLAG) $(INCLUDE) -M -MT $@ -o $(@:.o=.d) $<$(NVCC) $(CUFLAG) $(INCLUDE) -o $@ -c $<.PHONY: test
# PHONY是一个伪目标,它表示不管是否存在这个文件,只要执行这个目标,就会执行后面的命令
# 伪目标是指不生成任何文件,只是执行一些特定的命令
test:make cleanmakepython3 $(SOURCE_PY).PHONY: clean
clean:rm -rf ./*.d ./*.o ./*.so ./*.exe ./*.plan
加载使用
import ctypes
import osimport numpy as np
import tensorrt as trt
from cuda import cudartsoFile = "./AddScalarPlugin.so"
logger = trt.Logger(trt.Logger.ERROR)
trt.init_libnvinfer_plugins(logger, '')
# trt.init_libnvinfer_plugins函数的作用是初始化TensorRT库中的插件,其中的两个参数分别是日志级别和插件库的路径。
ctypes.cdll.LoadLibrary(soFile)
# ctypes.cdll.LoadLibrary函数的作用是加载指定的动态链接库,其中的参数是动态链接库的路径。
构建期
- TensorRT 向 Plugin 传输参数和权重
- Plugin 向 TensorRT 报告其输入输出张量信息,包括数量、形状(Shape)、数据类型(DataType)和数据排布(Layout)组合
- Plugin 向 TensorRT 报告其需要的 workspace 大小
- TensorRT 尝试各种允许的组合,选择性能最佳的输入输出组合(可能在 Plugin 前后插入 reformat 节点)
- Plugin 不参与层 fusing
def getAddScalarPlugin(scalar):for c in trt.get_plugin_registry().plugin_creator_list:#print(c.name)if c.name == "AddScalar":parameterList = []parameterList.append(trt.PluginField("scalar", np.float32(scalar), trt.PluginFieldType.FLOAT32))# PluginField类的作用是定义插件的属性,其中的三个参数分别是属性的名称、属性的值和属性的数据类型。return c.create_plugin(c.name, trt.PluginFieldCollection(parameterList))# create_plugin函数的作用是创建一个插件,其中的两个参数分别是插件的名称和插件的属性集合。return Nonebuilder = trt.Builder(logger)
network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))
profile = builder.create_optimization_profile()
config = builder.create_builder_config()inputT0 = network.add_input("inputT0", trt.float32, [-1 for i in shape])
profile.set_shape(inputT0.name, [1 for i in shape], [8 for i in shape], [32 for i in shape])
config.add_optimization_profile(profile)pluginLayer = network.add_plugin_v2([inputT0], getAddScalarPlugin(scalar))
# add_plugin_v2函数的作用是向网络中添加一个插件层,其中的两个参数分别是输入张量列表和插件。
network.mark_output(pluginLayer.get_output(0))
# mark_output函数的作用是标记网络的输出张量,其中的参数是张量。
engineString = builder.build_serialized_network(network, config)
# build_serialized_network函数的作用是构建序列化的网络,其中的两个参数分别是网络和配置。
运行期
- TensorRT 为 Plugin 提供输入输出张量的地址,workspace 的地址,以及所在的 stream
完整代码
import ctypes
import osimport numpy as np
import tensorrt as trt
from cuda import cudartsoFile = "./AddScalarPlugin.so"
np.set_printoptions(precision=3, linewidth=200, suppress=True)
# np.set_printoptions函数的作用是设置打印时的精度、行宽、是否使用科学计数法等。其中的三个
# 参数含义分别是:precision:设置浮点数的精度,即小数点后的位数;linewidth:设置输出的行宽;suppress:当suppress=True时,表示不输出小数点后面的数字,即将小数部分四舍五入
np.random.seed(31193)
cudart.cudaDeviceSynchronize()def printArrayInformation(x, info="", n=5):if 0 in x.shape:print('%s:%s' % (info, str(x.shape)))returnx = x.astype(np.float32)print( '%s:%s,SumAbs=%.5e,Var=%.5f,Max=%.5f,Min=%.5f,SAD=%.5f'%( \info,str(x.shape),np.sum(abs(x)),np.var(x),np.max(x),np.min(x),np.sum(np.abs(np.diff(x.reshape(-1)))) ))print('\t', x.reshape(-1)[:n], x.reshape(-1)[-n:])returndef check(a, b, weak=False, checkEpsilon=1e-5, info=""):if a.shape != b.shape:print("Error shape: A%s : B%s" % (str(a.shape), str(b.shape)))returnif weak:a = a.astype(np.float32)b = b.astype(np.float32)res = np.all(np.abs(a - b) < checkEpsilon)else:res = np.all(a == b)maxAbsDiff = np.max(np.abs(a - b))meanAbsDiff = np.mean(np.abs(a - b))maxRelDiff = np.max(np.abs(a - b) / (np.abs(b) + checkEpsilon))meanRelDiff = np.mean(np.abs(a - b) / (np.abs(b) + checkEpsilon))res = "%s:%s,MaxAbsDiff=%.2e,MeanAbsDiff=%.2e,MaxRelDiff=%.2e,MeanRelDiff=%.2e," % (info, res, maxAbsDiff, meanAbsDiff, maxRelDiff, meanRelDiff)index = np.argmax(np.abs(a - b))valueA, valueB= a.flatten()[index], b.flatten()[index]shape = a.shapeindexD = []for i in range(len(shape) - 1, -1, -1):x = index % shape[i]indexD = [x] + indexDindex = index // shape[i]res += "WorstPair=(%f:%f)at%s" %(valueA, valueB, str(indexD))print(res)returndef addScalarCPU(inputH, scalar):return [inputH[0] + scalar]def getAddScalarPlugin(scalar):for c in trt.get_plugin_registry().plugin_creator_list:#print(c.name)if c.name == "AddScalar":parameterList = []parameterList.append(trt.PluginField("scalar", np.float32(scalar), trt.PluginFieldType.FLOAT32))# PluginField类的作用是定义插件的属性,其中的三个参数分别是属性的名称、属性的值和属性的数据类型。return c.create_plugin(c.name, trt.PluginFieldCollection(parameterList))# create_plugin函数的作用是创建一个插件,其中的两个参数分别是插件的名称和插件的属性集合。return Nonedef run(shape, scalar):testCase = "<shape=%s,scalar=%f>" % (shape, scalar)trtFile = "./model-Dim%s.plan" % str(len(shape))print("Test %s" % testCase)logger = trt.Logger(trt.Logger.ERROR)trt.init_libnvinfer_plugins(logger, '')# trt.init_libnvinfer_plugins函数的作用是初始化TensorRT库中的插件,其中的两个参数分别是日志级别和插件库的路径。ctypes.cdll.LoadLibrary(soFile)# ctypes.cdll.LoadLibrary函数的作用是加载指定的动态链接库,其中的参数是动态链接库的路径。if os.path.isfile(trtFile):with open(trtFile, "rb") as f:engine = trt.Runtime(logger).deserialize_cuda_engine(f.read())if engine == None:print("Failed loading engine!")returnprint("Succeeded loading engine!")else:builder = trt.Builder(logger)network = builder.create_network(1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH))profile = builder.create_optimization_profile()config = builder.create_builder_config()inputT0 = network.add_input("inputT0", trt.float32, [-1 for i in shape])profile.set_shape(inputT0.name, [1 for i in shape], [8 for i in shape], [32 for i in shape])config.add_optimization_profile(profile)pluginLayer = network.add_plugin_v2([inputT0], getAddScalarPlugin(scalar))# add_plugin_v2函数的作用是向网络中添加一个插件层,其中的两个参数分别是输入张量列表和插件。network.mark_output(pluginLayer.get_output(0))# mark_output函数的作用是标记网络的输出张量,其中的参数是张量。engineString = builder.build_serialized_network(network, config)# build_serialized_network函数的作用是构建序列化的网络,其中的两个参数分别是网络和配置。if engineString == None:print("Failed building engine!")returnprint("Succeeded building engine!")with open(trtFile, "wb") as f:f.write(engineString)engine = trt.Runtime(logger).deserialize_cuda_engine(engineString)# deserialize_cuda_engine函数的作用是反序列化一个CUDA引擎,其中的参数是序列化的引擎。nIO = engine.num_io_tensors# num_io_tensors属性的作用是获取引擎的输入输出张量的数量。lTensorName = [engine.get_tensor_name(i) for i in range(nIO)]# get_tensor_name函数的作用是获取引擎的输入输出张量的名称。nInput = [engine.get_tensor_mode(lTensorName[i]) for i in range(nIO)].count(trt.TensorIOMode.INPUT)# get_tensor_mode函数的作用是获取引擎的输入输出张量的模式,其中的参数是张量的名称。context = engine.create_execution_context()context.set_input_shape(lTensorName[0], shape)#for i in range(nIO):# print("[%2d]%s->" % (i, "Input " if i < nInput else "Output"), engine.get_tensor_dtype(lTensorName[i]), engine.get_tensor_shape(lTensorName[i]), context.get_tensor_shape(lTensorName[i]), lTensorName[i])bufferH = []bufferH.append(np.arange(np.prod(shape), dtype=np.float32).reshape(shape))# np.arange函数的作用是创建一个等差数组,其中的参数是数组的大小。np.prod函数的作用是计算数组的元素个数。for i in range(nInput, nIO):bufferH.append(np.empty(context.get_tensor_shape(lTensorName[i]), dtype=trt.nptype(engine.get_tensor_dtype(lTensorName[i]))))# 初始化一个空数组,数组的形状是引擎的输入输出张量的形状,数组的数据类型是引擎的输出张量的数据类型。bufferD = []for i in range(nIO):bufferD.append(cudart.cudaMalloc(bufferH[i].nbytes)[1])# cudart.cudaMalloc函数的作用是在GPU上分配一块内存,其中的参数是内存的大小。# 为推理输入输出张量分配内存。for i in range(nInput):cudart.cudaMemcpy(bufferD[i], bufferH[i].ctypes.data, bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyHostToDevice)# cudart.cudaMemcpy函数的作用是在GPU之间复制内存,其中的四个参数分别是目标内存、源内存、内存的大小和复制的方向。# 将模型的输入张量从CPU复制到GPU。for i in range(nIO):context.set_tensor_address(lTensorName[i], int(bufferD[i]))# set_tensor_address函数的作用是设置张量的地址,其中的两个参数分别是张量的名称和地址。context.execute_async_v3(0)# execute_async_v3函数的作用是异步执行推理,其中的参数是批次大小。for i in range(nInput, nIO):cudart.cudaMemcpy(bufferH[i].ctypes.data, bufferD[i], bufferH[i].nbytes, cudart.cudaMemcpyKind.cudaMemcpyDeviceToHost)# 将模型的输出张量从GPU复制到CPU。outputCPU = addScalarCPU(bufferH[:nInput], scalar)"""for i in range(nInput):printArrayInformation(bufferH[i])for i in range(nInput, nIO):printArrayInformation(bufferH[i])for i in range(nInput, nIO):printArrayInformation(outputCPU[i - nInput])"""check(bufferH[nInput:][0], outputCPU[0], True)for b in bufferD:cudart.cudaFree(b)# 释放GPU上的内存。print("Test %s finish!\n" % testCase)if __name__ == "__main__":os.system("rm -rf ./*.plan")run([32], 1)run([32, 32], 1)run([16, 16, 16], 1)run([8, 8, 8, 8], 1)run([32], 1)run([32, 32], 1)run([16, 16, 16], 1)run([8, 8, 8, 8], 1)print("Test all finish!")