TensorRT的自定义算子Plugin的实现

这篇文章主要介绍了如何使用TensorRT实现自定义算子。

Note:

1. 我使用的是TensorRT7.0，自定义算子使用的IPluginV2IOExt实现的。
2. 模型框架是caffe，所以以下实现都只适用于caffe模型的解析，但理论上解析tf和onnx的改动不大。
3. 实现细节不方便全部贴出，但是基本实现过程和结构都在下面了，照着写写没啥问题了。

其实自定义算子写多了发现其实还挺好写的，格式都差不多，主要区别是enqueue的前向计算逻辑可能写起来复杂些。
整个实现过程基本上是：

1. 继承nvinfer1::IPluginV2IOExt，并实现相应的虚函数。
2. 继承nvinfer1::IPluginCreator并实现相应的虚函数。
3. 继承nvcaffeparser1::IPluginFactoryV2并实现相应的虚函数。
4. 在解析网络之前调用REGISTER_TENSORRT_PLUGIN注册UpsampleCreator和调用parser->setPluginFactoryV2()以使用自定义层类型。

以Upsample为例，TensorRT不支持Caffe的Upsample层，所以这里实现了一个自定义层类型，即plugin。需要实现：

1. Upsample类，继承自nvinfer1::IPluginV2IOExt。
2. UpsampleCreator类，继承自nvinfer1::IPluginCreator。
3. CaffePluginFactory类，继承自nvcaffeparser1::IPluginFactoryV2。

需要实现的函数详见如下代码段。

Upsample类的实现：

class Upsample : public nvinfer1::IPluginV2IOExt {
public:
    // 直接解析网络时候需要用到
    Upsample();
    // 反序列化时候需要用到
    Upsample(const void *data, size_t length);
    ~Upsample();
    
    // 直接return输出节点数，
    int getNbOutputs() override;
    
    // return输出的维度信息，如：return Dims3(inputs[0].d[0], inputs[0].d[1], inputs[0].d[2]);
    Dims getOutputDimensions(int index, const Dims *inputs, int num_input_dims) override;
    
    // pos索引到的input/output的数据格式（format）和数据类型（datatype）如果都支持则返回true
    bool supportsFormatCombination(int pos, const PluginTensorDesc* in_out, int num_inputs, int num_outputs) const override;
    
    // 这个函数可以获取到数据类型和输入的维度信息，如果有需要用到的可以在这里将相关信息取出来
    configurePlugin(const PluginTensorDesc* in, int num_inputs, const PluginTensorDesc* out, int num_outputs) override;

    // 在这里返回正确的序列化数据的长度，如我要序列化数据类型和数据维度：return sizeof(data_type) + sizeof(chw);
    size_t getSerializationSize() const override;
    
    // 序列化函数，在这里把反序列化时需要用到的参数或数据序列化
    void serialize(void *buffer) const override;
    
    // 设置工作空间，不需要直接 return 0;
    size_t getWorkspaceSize(int max_batch_size) const override;
    
    // 前向计算的核心函数，计算逻辑在这里实现，可以使用cublas实现或者自己写cuda核函数实现
    int enqueue(int batch_size, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
    
    // 调用enqueue的时候需要用到的资源先在这里Initialize，这个函数是在engine创建之后enqueue调用之前调用的，不需要Initialize则直接 return 0;
    int initialize() override;
    
    // 释放Initialize申请的资源，在enqueue调用之后且engine销毁之后调用
    void terminate() override;
    
    // 返回输出的数据类型，如何输入相同，可以直接 return input_types[0]；
    nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* input_types, int num_inputs) const override;
    
    // 返回自定义类型，如这里是：return Upsample
    const char* getPluginType() const override;
    
    // 返回plugin version，没啥说的
    const char* getPluginVersion() const override;
      
    // 销毁对象
    void destroy() override {
        delete this;
    }
    
    // 在这里new一个该自定义类型并返回
    nvinfer1::IPluginV2Ext* clone() const override;
    
    // 设置命名空间，用来在网络中查找和创建plugin
    void setPluginNamespace(const char* lib_namespace) override;
    // 返回plugin对象的命名空间
    const char* getPluginNamespace() const override;
    bool isOutputBroadcastAcrossBatch(int output_index, const bool* input_is_broadcasted, int num_inputs) const override;
    bool canBroadcastInputAcrossBatch(int input_index) const override;
}

下面是对应的Creator类的实现：

class UpsampleCreator : public nvinfer1::IPluginCreator {
public:
    const char* getPluginName() const override;
    const char* getPluginVersion() const override;
    const PluginFieldCollection* getFieldNames() override;
    // 创建自定义层pluin的对象并返回
    nvinfer1::IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) override;
    // 创建自定义层pluin的对象并返回，反序列化用到
    nvinfer1::IPluginV2* deserializePlugin(const char* name, const void* serial_data, size_t serial_length) override;
    void setPluginNamespace(const char* lib_namespace) override;
    const char* getPluginNamespace() const override;
}

下面是对应的plugin factory类的实现：

class CaffePluginFactory : public nvcaffeparser1::IPluginFactoryV2 {
public:
    // 在这里判断一个层是否为自定义层类型
    bool isPluginV2(const char* name) override;
    // 在这里创建自定义层类型的对象并返回
    nvinfer1::IPluginV2* createPlugin(const char* layer_name, const nvinfer1::Weights* weights, int num_weights, const char* libNamespace="") override;
}