深度学习自定义插件的编写

实现自定义插件的意义

        TensorRT在使用过程中会遇到一些问题，如算子不支持，自动生成的算子结构冗余，对fp16，int8的精度误差放大等问题，这种时候可以考虑自己实现相关算子，同时使用onnx-graphsurgen对onnx模型进行编辑，包括节点的增加、删除、连接，自己修改的模型很可能出现结果不正确的情况，可以配合使用polygraphy对模型精度做调试，onnx-graphsurgen是英伟达推出的对onnx模型做编辑的工具，polygraphy也是英伟达推出的onnx模型调试工具，能同时运行修改前和修改后的onnx模型，并且支持逐层运行以方便精度对齐。

构建含有onnx插件的模型

        自定义插件的构成包括模型训练时的构建，导出onnx，再到对应部署平台的转化三个部分，本文将举例switsh运算模块的pytorch模块构建与使用、导出onnx，最后使用TensorRT的插件模块完成注册的全流程。

完整代码见/saveDemo/Plugin  learn: 储存相关demo

首先使用pytorch实现switsh算子，要包括前向传播模块、反向传播模块、onnx导出模块三部分

class mySwish(torch.autograd.Function):
    @staticmethod
    def symbolic(g, input, bias): #用于导出onnx、
        # g表示用于构建onnx图，g.op表示返回一个操作节点，名称为"Plugin"，输入为input、bias，
        # 其中name_s为该节点添加名字描述，infos同样用于信息描述
        return g.op("custom::Plugin", input, bias, name_s="Swish", info_s=json.dumps({
            "size": 555,
            "shape": [6, 7, 8],
            "module": "abcdefg"
        }))

    @staticmethod
    def forward(ctx, input, bias): #输入包括上下文(pytorch自动维护)权重和偏置
        ctx.save_for_backward(input)# 保存输入用于反向传播
        return input * torch.sigmoid(input) + bias # 正向传播表达式

    @staticmethod
    def backward(ctx, grad_output):
        input = ctx.saved_tensors[0]# 获取保存的输入
        sigmoid_input = torch.sigmoid(input) 
        return grad_output * (sigmoid_input * (1 - sigmoid_input)), grad_output # 返回值包括权重和偏置
        # 返回值为链式法则求导，其中grad_output为loss对结果求导的值，再乘以结果对x的导数

这里反向传播代码的校验pytorch提供了方法，参考如下：

input_tensor = torch.ones(3, requires_grad=True, dtype=torch.double)
bias = torch.ones(3, requires_grad=True, dtype=torch.double)
#eps为极小值，防止做除数，atol表示对误差的容忍度
input_grad_check = gradcheck(mySigmoid.apply, (input_tensor, bias), eps=1e-5, atol=1e-4)
bias_grad_check = gradcheck(mySigmoid.apply, (input_tensor, bias), eps=1e-5, atol=1e-4)
print("Numerical gradient check for input:", input_grad_check)
print("Numerical gradient check for bias:", bias_grad_check)

结果正确时会有如下输出，其余结果均表示有问题

Numerical gradient check for input: True
Numerical gradient check for bias: True

接下来在构建模型时使用自定义的算子，由于实现时使用了静态方法，因此使用模块时只需要apply就行，使用下面的模块给调用时添加偏置项

class MemoryEfficientSwish(nn.Module):
    def __init__(self):
        super().__init__()
        
        # 这里我们假设有bias作为权重参数
        self.bias = nn.Parameter(torch.zeros(1))
        self.bias.data.fill_(3.15)

    def forward(self, x):
        # 我们假设丢一个bias进去
        return mySwish.apply(x, self.bias)

写一个简单的3*3的卷积神经网络调用刚刚写的自定义模块：

class Model(nn.Module):
    def __init__(self):
        super().__init__()

        self.conv = nn.Conv2d(1, 1, 3, stride=1, padding=0, dilation=1, bias=False)
        self.conv.weight.data = torch.FloatTensor([
            [1, 0, 0],
            [0, 1, 0],
            [0, 0, 1]
        ]).view(1, 1, 3, 3)
        self.swish = MemoryEfficientSwish()

    def forward(self, x):
        return self.swish(self.conv(x))

然后就可以将模型导出onnx准备部署，需要将自定义模块注册到torch.onnx中才能导出，此步骤需要注意2个点，首先是注册模块的命名规则为作用域::插件名，类似于c++的命名空间，其次是注册模块的opset_version要设置为1，设置为11会莫名报错。

domain_name = "custom" # 设置命名空间
torch.onnx.register_custom_op_symbolic(f"{domain_name}::Plugin", mySwish.symbolic, 1) # 注册插件名到torch.onnx
torch.onnx.export(model, (input,), "plugin.onnx", #正常导出
    input_names=["Input"], 
    output_names=["Output"],
    opset_version=11,
    verbose=True)

实现Tensorrt的插件

        现在训练好模型并导出onnx之后就可以使用相关框架做模型部署，下面将以TensorRT的c++接口举例自定义模块的使用，在TensorRT中无法直接调用自己定义的pytorch算子，通常的做法是将算子用TensorRT提供的接口实现后打包成动态库，在导出、推理时链接相关库，主要操作步骤包括：1.编写plugin，2.实现插件创建类，3.注册插件到TRT库，4.序列化插件，反序列化过程包括：1.获取plugin的类型和版本，2.获取creater，3.使用creater生成plugin并插入到网络中。

plugin的构造函数：

        network define时、clone时、creater被调用时会触发plugin的构造函数，其中define时调用的构造函数必须手动创建，clone函数选填，实现这两个基本就满足使用需求了；

        define时调用的构造函数参数比较灵活，可根据情况自行定义，反序列化时的参数比较固定，仅包含数据地址和长度。

plugin相关api：

        int getNbOutputs() const; //获得输出的个数

        nvinfer1::Dims getOutputDimensions(int index, const nvinfer1::Dims* inputs, int nbInputDims);//根据index获取第index个输出的维度

        nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const;//根据输入和index获取第index个输出的类型

        size_t getSerializationSize() const;//获取参数、权重等占内存的大小

        void serialize(void* buffer) const;// 将参数、权重写入到buffer中

        const char* getPluginType() const;// 获取plugin名称，必须唯一

        const char* getPluginVersion() const; //获取plugin版本，必须唯一

        int initialize(); //申请权值显存空间并拷贝

        void terminate(); //释放initialize开辟的空间

        void destroy(); //释放整个plugin占用的资源，terminate写好就够了

        void configurePlugin(const nvinfer1::PluginTensorDesc* in, int nbInput, const nvinfer1::PluginTensorDesc* out, int nbOutput); //判断输入输出数据类型、shape等是否正确，以及运行时的参数初始化

        bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int bnInputs, int bnOutputs) const; //判断输入输出的数据类型

        size_t getWorkspaceSize(int maxBatchSize) const;// 查询工作空间

        //工作空间的使用可以防止单个模型的显存溢出，以及显存复用

        int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream);//执行模型，batchSize可以小于MaxBatchSize，workspace自动根据batchsize做调整，

        nvinfer1::IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength); //反序列化，生成plugin

如果要添加的模块在模型的前处理或者后处理部分则不需要使用torch实现，而是直接在onnx中添加相关节点描述，而后在trt中实现即可，tensorrt的算子实现之前已给出，onnx相关操作方法将在下方文章给出

        onnx节点包含op(操作类型)，name(节点名称)，attrs(字典信息)，inputs(输入Tensor信息，list)，outputs(输出节点信息，list)，onnx中的Tensor包含三个属性：name、dtype、shape，其中shape通常仅输入节点是运行前已知的，过程节点和输出的shape只有在运行之后才能知道。编写onnx节点时如果不含权重信息（如sigmoid等激活函数）则只需要定义节点的输入输出信息，包括名称和shape、dtype等，如果包含权重信息则需要将额外的信息写入到attrs中，

        完整的demo请参照 /saveDemo/Plugin  learn: 储存相关demo，已经逐行加注释，有问题欢迎留言。