TensorRT plugin插件编写

TensorRT plugin插件编写

---

参考

CUDA与TensorRT部署部署实战第四章
实现TensorRT自定义插件(plugin)自由！
TensorRT教程笔记–Plugin

学习目标

• TensorRT自定义插件的使用方式
• 如何添加自己的自定义算子

注意：本文中主要用来对函数的作用和意义进行说明，故参考函数的编写并非只针对一个插件。

---

引言

1. 下图中TensorRT不同版本对应的plugin基类，推荐IPluginV2DynamicExt， 版本新，支持动态shape。
   

2. TensorRT 官方插件样例，地址：https://github.com/NVIDIA/TensorRT/tree/master/plugin
   内部包含插件的源码。
   如果要添加自己的算子，可以在官方的plugin库里头进行修改添加，然后编译官方的plugin库。将生成的libnvinfer_plugin.so.7替换原本的.so文件即可。

一. 导出自定义op的onnx模型

参考：onnx中注册算子的方法
Eg. 使用方式三注册算子类，添加两个静态方法：一个符号函数symbolic和一个forward。

class CustomScalarImpl(torch.autograd.Function):
    @staticmethod
    ## g: 代表torch.Graph, x 代表输入tensor,  其余为传入参数
    def symbolic(g, x, r, s):
       ## 额外添加参数可以使用 xx_f （浮点型）, xx_i （整形）的形式
        return g.op("custom::customScalar", x, scalar_f=r, scale_f=s) 
    @staticmethod
    ## ctx 代表上下文参数
    def forward(ctx, x, r, s):
        return (x + r) * s

---

二. 自定义插件流程

我们需要写两个类：
MyCustomPlugin，继承IPluginV2DynamicExt，是插件类，用于写插件具体的实现，最主要函数enqueue。
MyCustomPluginCreator，继承IPluginCreator，是插件工厂类，用于根据需求创建该插件。

2.1 MyCustomPlugin 插件类

1. 总览

class CustomLeakyReLUPlugin : public IPluginV2DynamicExt {
public:
    CustomLeakyReLUPlugin() = delete; 
    CustomLeakyReLUPlugin(const std::string &name, float alpha);
    CustomLeakyReLUPlugin(const std::string &name, const void* buffer, size_t length);

    ~CustomLeakyReLUPlugin();

    const char* getPluginType() const noexcept override;
    const char* getPluginVersion() const noexcept override;
    int32_t     getNbOutputs() const noexcept override;
    size_t      getSerializationSize() const noexcept override;
    const char* getPluginNamespace() const noexcept override;
    DataType    getOutputDataType(int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
    DimsExprs   getOutputDimensions(int32_t outputIndex, const DimsExprs* input, int32_t nbInputs, IExprBuilder &exprBuilder) noexcept override;
    size_t      getWorkspaceSize(const PluginTensorDesc *inputs, int32_t nbInputs, const PluginTensorDesc *outputs, int32_t nbOutputs) const noexcept override;

    int32_t     initialize() noexcept override;
    void        terminate() noexcept override;
    void        serialize(void *buffer) const noexcept override;
    void        destroy() noexcept override;
    int32_t     enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, const void* const* ionputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override; // 实际插件op执行的地方，具体实现forward的推理的CUDA/C++实现会放在这里面
    IPluginV2DynamicExt* clone() const noexcept override;

    bool        supportsFormatCombination(int32_t pos, const PluginTensorDesc* inOuts, int32_t nbInputs, int32_t nbOutputs) noexcept override;
    void        configurePlugin(const DynamicPluginTensorDesc* in, int32_t nbInputs, const DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept override;
    void        setPluginNamespace(const char* pluginNamespace) noexcept override;

    void        attachToContext(cudnnContext* contextCudnn, cublasContext* contextCublas, IGpuAllocator *gpuAllocator) noexcept override;
    void        detachFromContext() noexcept override;

private:
    const std::string mName;
    std::string       mNamespace;
    struct {
        float alpha;
    } mParams;  // 可以使用结构体的形式来封装参数
};

2. 成员变量
如果你的插件有weights(类似于conv操作的weight和bias)，有参数(类似于conv中的kernel-size、padding)，在类中则需要定义为成员变量，为private类型。
Eg.

private:
    std::vector _weight; // 权重，在cpu空间存放
    std::vector _bias;   // 权重偏置，在cpu空间存放
    float* _d_weight;           // 权重，在GPU空间存放
    float* _d_bias;				// 权重偏置，在GPU空间存放
    bool _initialized;
    std::string mNamespace;   // 命名空间
    const std::string mName;   // 插件名字

3. 构造函数与析构函数
三个构造函数，在三个阶段被调用： parse, clone, 反序列化。

    /* 默认构造函数，删除 */
    CustomLeakyReLUPlugin() = delete; 
    
    /* parse阶段： 
    	得到onnx模型后，parse这个插件，会读取参数信息转换为TensorRT格式；
    	PluginCreator用于创建该插件时调用的构造函数*/
    CustomLeakyReLUPlugin(const std::string &name, float alpha);
    
    /* clone和deseriaze阶段：
    	clone阶段在parse之后，TensorRT为了优化插件，生成许多优化版本去测试，也可以在推理阶段供不同的context创建插件使用。
    	将序列化好的Plugin进行反序列化的时候也需要创建插件的实例 */
    CustomLeakyReLUPlugin(const std::string &name, const void* buffer, size_t length);

析构函数则需要执行terminate和destroy，terminate函数就是释放这个op之前开辟的一些显存空间。

4. 成员函数
1）有关获取plugin信息的方法

# 获取输出tensor的数目， 一般而言是1
int32_t CustomLeakyReLUPlugin::getNbOutputs() const noexcept
{
    return 1;
}

# 获取输出数据的类型，通常与输入数据tensor保持一致
DataType CustomLeakyReLUPlugin::getOutputDataType(int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
{
	ASSERT(inputTypes && nbInputs > 0 && index == 0);
    return inputTypes[0];
}

# 获取中间显存变量的实际数据大小(`bytesize`)
size_t MyCustomPlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const 
{ 
    // 计算这个op前向过程中你认为需要的中间显存数量
    size_t need_num;
    return need_num * sizeof(float);
}

# 这个成员函数中根据输入维度推理出模型的输出维度，需要注意的是，但这个输出维度其实“内定”的(也就是在计算之前就算出来了)
DimsExprs CustomLeakyReLUPlugin::getOutputDimensions(int32_t outputIndex, const DimsExprs* inputs, int32_t nbInputs, IExprBuilder &exprBuilder) noexcept
{
    return inputs[0];
}

# set/getPluginNamespace
为这个插件设置namespace名字，如果不设置则默认是""，需要注意的是同一个namespace下的plugin如果名字相同会冲突。
const char* CustomLeakyReLUPlugin::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

# getSerializationSize： 返回序列化的字节数
size_t CustomLeakyReLUPlugin::getSerializationSize() const noexcept
{
    return sizeof(mParams);
}

2）实际op执行函数： enqueue

int32_t CustomLeakyReLUPlugin::enqueue(
    const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, 
    const void* const* inputs, void* const* outputs, 
    void* workspace, cudaStream_t stream) noexcept
{
    int nElements = 1;
    for (int i = 0; i < inputDesc[0].dims.nbDims; i++){
        nElements *= inputDesc[0].dims.d[i];
    }

    customLeakyReLUImpl(
            static_cast(inputs[0]),
            static_cast(outputs[0]), 
            mParams.alpha, 
            nElements,
            stream);

    return 0;
}

Tips.
默认写的.cu是fp32的，TensorRT在fp16运行模式下，运行到不支持fp16的插件op时，会自动切换到fp32模式，等插件op运行完再切换回来。

3）configurePlugin
配置这个插件op，判断输入和输出类型数量是否正确。官方还提到通过这个配置信息可以告知TensorRT去选择合适的算法(algorithm)去调优这个模型。

void MyCustomPluginDynamic::configurePlugin(
    const nvinfer1::DynamicPluginTensorDesc *inputs, int nbInputs,
    const nvinfer1::DynamicPluginTensorDesc *outputs, int nbOutputs) {
  // Validate input arguments
  assert(nbOutputs == 1);
  assert(nbInputs == 2);
  assert(mType == inputs[0].desc.type);
}

4）clone
调用第三个构造函数，进行克隆，clone成员函数主要用于传递不变的权重和参数，将plugin复制n多份，从而可以被不同engine或者builder或者network使用。

IPluginV2DynamicExt* CustomLeakyReLUPlugin::clone() const noexcept
{
    try{
        auto p = new CustomLeakyReLUPlugin(mName, &mParams, sizeof(mParams));
        p->setPluginNamespace(mNamespace.c_str());
        return p;
    }
    catch (std::exception const &e){
        LOGE("ERROR detected when clone plugin: %s", e.what());
    }
    return nullptr;
}

5）supportsFormatCombination
TensorRT调用此方法以判断pos索引的输入/输出是否支持inOut[pos].format和inOut[pos].type指定的格式/数据类型。

bool CustomLeakyReLUPlugin::supportsFormatCombination(int32_t pos, const PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
    switch (pos) {
    case 0:
        return inOut[0].type == DataType::kFLOAT && inOut[0].format == TensorFormat::kLINEAR;
    case 1:
        return inOut[1].type == DataType::kFLOAT && inOut[1].format == TensorFormat::kLINEAR;
    default:
        return false;
    }
    return false;
}

6） serialize
把需要用的数据按照顺序序列化到buffer里头。

void CustomLeakyReLUPlugin::serialize(void *buffer) const noexcept
{
    memcpy(buffer, &mParams, sizeof(mParams));
    return;

}

7）attachToContext
如果这个op使用到了一些其他东西，例如cublas handle，可以直接借助TensorRT内部提供的cublas handle。

void MyCustomPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator)
{
     mCublas = cublasContext;
}

2.2 MyCustomPluginCreator 插件工厂类

总览

class MyCustomPluginCreator : public BaseCreator
{
public:
  MyCustomPluginCreator(); //初始化mFC以及mAttrs
  ~MyCustomPluginCreator() override = default;
  const char* getPluginName() const override;    
  const char* getPluginVersion() const override; 
  const PluginFieldCollection* getFieldNames() override;
  IPluginV2DynamicExt* createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc) override;
  IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
private:
  static PluginFieldCollection mFC; //接受plugionFields传进来的权重和参数，并将信息传递给Plugin，内部通过createPlugin来创建带参数的plugin
  static std::vector mAttrs;  //用来保存这个插件op所需要的权重和参数,
                                           //从onnx中获取, 同样在`parse`的时候使用
  std::string mNamespace;
};

1）构造函数和析构函数

CustomScalarPluginCreator::CustomScalarPluginCreator()
{
    /* 
     * 每个插件的Creator构造函数需要定制，主要就是获取参数以及传递参数
     * 初始化creator中的PluginField以及PluginFieldCollection
     * - PluginField::            负责获取onnx中的参数
     * - PluginFieldCollection：  负责将onnx中的参数传递给Plugin， 内部通过createPlugin来创建带参数的plugin
    */

    mAttrs.emplace_back(PluginField("scalar", nullptr, PluginFieldType::kFLOAT32, 1));
    mAttrs.emplace_back(PluginField("scale", nullptr, PluginFieldType::kFLOAT32, 1));
    mFC.nbFields = mAttrs.size();
    mFC.fields   = mAttrs.data();
}

2）createPlugin
这个成员函数作用是通过PluginFieldCollection去创建plugin，将op需要的权重和参数一个一个取出来，然后调用上文提到的第一个构造函数。

IPluginV2DynamicExt* MyCustomPlugin::createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc)
{
    int in_channel;
    std::vector weight;
    std::vector bias;
    const PluginField* fields = fc->fields;
    for (int i = 0; i < fc->nbFields; ++i)
    {
        const char* attrName = fields[i].name;
        if (!strcmp(attrName, "in_channel"))
        {
            ASSERT(fields[i].type == PluginFieldType::kINT32);
            in_channel= *(static_cast(fields[i].data));
        }
        else if (!strcmp(attrName, "weight"))
        {
            ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
            int size = fields[i].length;
            h_weight.reserve(size);
            const auto* w = static_cast(fields[i].data);
            for (int j = 0; j < size; j++)
            {
                h_weight.push_back(*w);
                w++;
            }
        }
        else if (!strcmp(attrName, "bias"))
        {
            ASSERT(fields[i].type == PluginFieldType::kFLOAT32);
            int size = fields[i].length;
            h_bias.reserve(size);
            const auto* w = static_cast(fields[i].data);
            for (int j = 0; j < size; j++)
            {
                h_bias.push_back(*w);
                w++;
            }
        }
    }

    Weights weightWeights{DataType::kFLOAT, weight.data(), (int64_t) weight.size()};
    Weights biasWeights{DataType::kFLOAT, bias.data(), (int64_t)_bias.size()};

    MyCustomPlugin* obj = new MyCustomPlugin(in_channel, weightWeights, biasWeights);
    obj->setPluginNamespace(mNamespace.c_str());
    return obj;
}

3）deserializePlugin
这个函数会被onnx-tensorrt的一个叫做TRT_PluginV2的转换op调用，这个op会读取onnx模型的data数据将其反序列化到network中。

IPluginV2* CustomLeakyReLUPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept
{
    try{
    # new 一个plugin
        return new CustomLeakyReLUPlugin(name, serialData, serialLength);
    }
    catch (std::exception const &e){
        LOGE("ERROR detected when deserialize plugin: %s", e.what());
    }
    return nullptr;
}

4）成员变量PluginFieldCollection
这个是成员变量，也会作为getFieldNames成员函数的返回类型。PluginFieldCollection的主要作用是传递这个插件op所需要的权重和参数给到createPlugin使用，在实际的engine推理过程中并不使用，而在parse中会用到(例如caffe2trt、onnx2trt)。

当使用这些parse去解析这个op的时候，这个op的权重和参数会经历Models --> TensorRT engine --> TensorRT runtime这个过程。

---

三. 注册插件

主要有两种方法：
第一种可以看官方的plugin代码。

extern "C" {
bool initLibNvInferPlugins(void* logger, const char* libNamespace)
{
    initializePlugin(logger, libNamespace);
    initializePlugin(logger, libNamespace);
    initializePlugin(logger, libNamespace);
    ...
    return true;
}
其中initializePlugin函数执行了addPluginCreator函数：

template 
void initializePlugin(void* logger, const char* libNamespace)
{
    PluginCreatorRegistry::getInstance().addPluginCreator(logger, libNamespace);
}

addPluginCreator函数又执行了getPluginRegistry()->registerCreator对pluginCreator进行了注册，这样就完成注册任务了。

第二种，很简单，在实现文件的custom命名空间内添加如下，通过宏的方式注册。

REGISTER_TENSORRT_PLUGIN(CustomScalarPluginCreator);