深度学习模型部署(十一)TensorRT写Plugin

最新推荐文章于 2024-09-02 19:54:56 发布

战术摸鱼大师

最新推荐文章于 2024-09-02 19:54:56 发布

阅读量688

点赞数 6

分类专栏：模型部署实战文章标签：深度学习人工智能神经网络

本文链接：https://blog.csdn.net/qq_45983373/article/details/136771662

版权

模型部署实战专栏收录该内容

19 篇文章 11 订阅

订阅专栏

什么是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 IPluginV2
    const 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 IPluginV2Ext
    DataType 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 IPluginV2DynamicExt
    IPluginV2DynamicExt *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.inc

SOURCE_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 clean
	make
	python3 $(SOURCE_PY)

.PHONY: clean
clean:
	rm -rf ./*.d ./*.o ./*.so ./*.exe ./*.plan

加载使用

import ctypes
import os

import numpy as np
import tensorrt as trt
from cuda import cudart

soFile = "./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 None

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函数的作用是构建序列化的网络，其中的两个参数分别是网络和配置。

运行期

TensorRT 为 Plugin 提供输入输出张量的地址，workspace 的地址，以及所在的 stream

完整代码

import ctypes
import os

import numpy as np
import tensorrt as trt
from cuda import cudart

soFile = "./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)))
        return
    x = 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:])
    return

def 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)))
        return
    if 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.shape
    indexD = []
    for i in range(len(shape) - 1, -1, -1):
        x = index % shape[i]
        indexD = [x] + indexD
        index = index // shape[i]
    res += "WorstPair=(%f:%f)at%s" %(valueA, valueB, str(indexD))
    print(res)
    return

def 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 None

def 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!")
            return
        print("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!")
            return
        print("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!")