5.4.tensorRT基础(2)-学习第一个插件的编写

前言

杜老师推出的 tensorRT从零起步高性能部署 课程，之前有看过一遍，但是没有做笔记，很多东西也忘了。这次重新撸一遍，顺便记记笔记。

本次课程学习 tensorRT 基础-学习第一个插件的编写

课程大纲可看下面的思维导图

1. 插件

这节我们学习插件的实现

插件实现-重点：

1. 如何在 pytorch 里面导出一个插件

2. 插件解析时如何对应，在 onnx parser 中如何处理

3. 插件的 creator 实现

4. 插件的具体实现，继承自 IPluginV2DynamicExt

5. 插件的序列化和反序列化

首先我们提供了一个 gen-onnx.py 文件，用于在 pytorch 中导出一个插件，内容如下：

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.onnx
import torch.autograd
import os

class MYSELUImpl(torch.autograd.Function):

    # reference: https://pytorch.org/docs/1.10/onnx.html#torch-autograd-functions
    @staticmethod
    def symbolic(g, x, p):
        print("==================================call symbolic")
        return g.op("MYSELU", x, p, 
            g.op("Constant", value_t=torch.tensor([3, 2, 1], dtype=torch.float32)),
            attr1_s="这是字符串属性", 
            attr2_i=[1, 2, 3], 
            attr3_f=222
        )

    @staticmethod
    def forward(ctx, x, p):
        return x * 1 / (1 + torch.exp(-x))


class MYSELU(nn.Module):
    def __init__(self, n):
        super().__init__()
        self.param = nn.parameter.Parameter(torch.arange(n).float())

    def forward(self, x):
        return MYSELUImpl.apply(x, self.param)


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

        self.conv = nn.Conv2d(1, 1, 3, padding=1)
        self.myselu = MYSELU(3)
        self.conv.weight.data.fill_(1)
        self.conv.bias.data.fill_(0)
    
    def forward(self, x):
        x = self.conv(x)
        x = self.myselu(x)
        return x


# 这个包对应opset11的导出代码，如果想修改导出的细节，可以在这里修改代码
# import torch.onnx.symbolic_opset11
print("对应opset文件夹代码在这里：", os.path.dirname(torch.onnx.__file__))

model = Model().eval()
input = torch.tensor([
    # batch 0
    [
        [1,   1,   1],
        [1,   1,   1],
        [1,   1,   1],
    ],
        # batch 1
    [
        [-1,   1,   1],
        [1,   0,   1],
        [1,   1,   -1]
    ]
], dtype=torch.float32).view(2, 1, 3, 3)

output = model(input)
print(f"inference output = \n{output}")

dummy = torch.zeros(1, 1, 3, 3)
torch.onnx.export(
    model, 

    # 这里的args，是指输入给model的参数，需要传递tuple，因此用括号
    (dummy,), 

    # 储存的文件路径
    "workspace/demo.onnx", 

    # 打印详细信息
    verbose=True, 

    # 为输入和输出节点指定名称，方便后面查看或者操作
    input_names=["image"], 
    output_names=["output"], 

    # 这里的opset，指，各类算子以何种方式导出，对应于symbolic_opset11
    opset_version=11, 

    # 表示他有batch、height、width3个维度是动态的，在onnx中给其赋值为-1
    # 通常，我们只设置batch为动态，其他的避免动态
    dynamic_axes={
        "image": {0: "batch", 2: "height", 3: "width"},
        "output": {0: "batch", 2: "height", 3: "width"},
    },

    # 对于插件，需要禁用onnx检查
    enable_onnx_checker=False
)

print("Done.!")

其中 Model 类还是和之前一样的 conv+activation，不过这次激活函数替换成了我们自己实现的一个 MYSELU，其 forward 是调用了 MYSELUImpl 这个类，它是继承自 torch.autograd.function，只要是插件你就要继承自它，在这个插件实现类中你需要实现一个静态 forward 方法，它接收 ctx，x，param 三个参数。

你想实现自定义算子的话，你可以自己实现 forward 甚至是 backward，另外 forwad 中的部分不会被跟踪，也就是意味着不会生成节点，那真正生成节点靠的是 symbolic 这个静态函数，具体可参考 pytorch 官方文档的 https://pytorch.org/docs/1.10/onnx.html#torch-autograd-functions

它接收 g，x，param 三个参数，通过 g.op 创建一个节点，我们来导出一下 onnx，出现了如下错误：

图1-1 enable_onnx_checker错误

提示 onnx.export() 没有 enable_onnx_checker 这个参数，应该是我的 torch 版本的原因，换了我的笔记本电脑没问题

图1-2 onnx正常导出

导出完成之后还打印了推理结果，待会和 C++ 的推理结果做对比，导出的 onnx 如下图所示：

图1-3 demo.onnx

可以看到激活函数是我们自定义实现的 MYSELU，它还有 Constant 这个节点输出，是一些常量值，接下来我们执行下 C++ 的代码来解析下这个 onnx 文件，看能不能拿到正确的结果，执行下 make run

图1-4 make run结果

可以看到正常解析并执行完推理了，同时推理结果和 pytorch 保持一致，说明我们的插件能够被正常执行，main.cpp 中除了包含自己编译的 onnx-tensorrt-release-8.0 之外没有任何不同，那 MYSELU 这个算子是如何被解析的呢？

在 builtin_op_importers.cpp 最后我们注册了一个叫做 MYSELU 的节点，内容如下：

DEFINE_BUILTIN_OP_IMPORTER(MYSELU)
{
    printf("\033[31m=======================call MYSELU==============\033[0m\n");
    OnnxAttrs attrs(node, ctx);
    const std::string pluginName{node.op_type()};
    const std::string pluginVersion{attrs.get<std::string>("plugin_version", "1")};
    const std::string pluginNamespace{attrs.get<std::string>("plugin_namespace", "")};

    LOG_INFO("Searching for plugin: " << pluginName << ", plugin_version: " << pluginVersion << ", plugin_namespace: " << pluginNamespace);
    nvinfer1::IPluginCreator* creator = importPluginCreator(pluginName, pluginVersion, pluginNamespace);
    ASSERT(creator && "Plugin not found, are the plugin name, version, and namespace correct?", ErrorCode::kUNSUPPORTED_NODE);

    const nvinfer1::PluginFieldCollection* fieldNames = creator->getFieldNames();
    // Field data needs to be type erased, we use fieldData for temporary allocations.
    string_map<std::vector<uint8_t>> fieldData{};
    std::vector<nvinfer1::PluginField> fields = loadFields(fieldData, attrs, fieldNames, ctx);

    const auto plugin = createPlugin(getNodeName(node), creator, fields);
    ASSERT(plugin && "Could not create plugin", ErrorCode::kUNSUPPORTED_NODE);

    std::vector<nvinfer1::ITensor*> pluginInputs{};
    for (auto& input : inputs)
    {
        pluginInputs.emplace_back(&convertToTensor(input, ctx));
    }
    LOG_INFO("Successfully created plugin: " << pluginName);
    auto* layer = ctx->network()->addPluginV2(pluginInputs.data(), pluginInputs.size(), *plugin);
    ctx->registerLayer(layer, getNodeName(node));
    RETURN_ALL_OUTPUTS(layer);
}

正是这个节点的注册才能让 onnx 解析器正常解析这个节点，那这个节点解析出来了怎么对应到执行的呢，我们通过 importPluginCreator() 将 pluginName 即 node.op_type() 即 MYSELU 当参数传入创建一个 creator 得到一个插件的创建器，通过 creator 去读取 onnx 对应的属性，然后将对应的属性丢到 createPlugin() 这个函数去创建我们的 plugin，然后通过 ctx->network->addPluginV2() 得到对应的 layer，最后通过 registerLayer 完成我们插件的注册。

那这个只是插件的创建和注册，并没有说明它是怎么和我们的执行代码关联起来的，我们来看下 myselu-plugin.hpp 和 myselu-plugin.cpp 这两个文件，首先头文件它要实现 PluginCreater，要实现获取插件的名字、版本号、创建插件实例、反序列化插件…，内容如下：

#ifndef CUSTOM_MYSELU_PLUGIN_H
#define CUSTOM_MYSELU_PLUGIN_H

#include "NvInferPlugin.h"
#include <string>
#include <vector>

using namespace nvinfer1;

class MySELUPlugin : public IPluginV2DynamicExt
{
public:
    MySELUPlugin(const std::string name, const std::string attr1, float attr3);

    MySELUPlugin(const std::string name, const void* data, size_t length);

    // It doesn't make sense to make MySELUPlugin without arguments, so we delete default constructor.
    MySELUPlugin() = delete;

    int getNbOutputs() const noexcept override;

    virtual nvinfer1::DataType getOutputDataType(
        int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept override{
        return inputTypes[0];
    }

    virtual nvinfer1::DimsExprs getOutputDimensions(
        	int32_t outputIndex, const nvinfer1::DimsExprs* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept override;

    int initialize() noexcept override;

    void terminate() noexcept override;

    virtual size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int32_t nbInputs, const nvinfer1::PluginTensorDesc* outputs,
        	int32_t nbOutputs) const noexcept override{
        return 0;
    }

    int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
            const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;

    size_t getSerializationSize() const noexcept override;

    void serialize(void* buffer) const noexcept override;

    virtual void configurePlugin(const DynamicPluginTensorDesc* in, int32_t nbInputs,
        const DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept;

    virtual bool supportsFormatCombination(int32_t pos, const PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;

    const char* getPluginType() const noexcept override;

    const char* getPluginVersion() const noexcept override;

    void destroy() noexcept override;

    nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;

    void setPluginNamespace(const char* pluginNamespace) noexcept override;

    const char* getPluginNamespace() const noexcept override;

private:
    const std::string mLayerName;
    std::string mattr1;
    float mattr3;
    size_t mInputVolume;
    std::string mNamespace;
};

class MySELUPluginCreator : public IPluginCreator
{
public:
    MySELUPluginCreator();

    const char* getPluginName() const noexcept override;

    const char* getPluginVersion() const noexcept override;

    const PluginFieldCollection* getFieldNames() noexcept override;

    IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override;

    IPluginV2* 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;

private:
    static PluginFieldCollection mFC;
    static std::vector<PluginField> mPluginAttributes;
    std::string mNamespace;
};

#endif

创建插件实例将参数序列化为文件，不再走 onnx，因此反序列化你也需要实现。其次头文件中还有插件类的实现，这是我们关注的重点，首先它继承自 IPluginV2DynamicExt，我们要实现两种构造方式，通过参数来构造以及通过数据来构造，然后你还需要实现一些特定的函数，比如获取输出的数量、输出的数据类型、输出的维度…，其中 enqueue 才是我们执行时使用的函数，它不仅仅在运行时才执行，在编译阶段 tensorRT 会以实际方式在设备上运行以此选择哪个配置是最合适的。

对应的实现在 CPP 文件中，其内容如下：

#include "myselu-plugin.hpp"
#include "NvInfer.h"

#include <cassert>
#include <cstring>
#include <vector>

using namespace nvinfer1;

void myselu_inference(const float* x, float* output, int n, cudaStream_t stream);

// MySELU plugin的特定常量
namespace
{
const char* MYSELU_PLUGIN_VERSION{"1"};  // 采用的名称要对应上onnx-tensorrt-release-8.0/builtin_op_importers.cpp:5094行定义的名称
const char* MYSELU_PLUGIN_NAME{"MYSELU"};
} // namespace

// 静态类字段的初始化
PluginFieldCollection MySELUPluginCreator::mFC{}; // FieldCollection 字段收集
std::vector<PluginField> MySELUPluginCreator::mPluginAttributes;
// 实际注册时，注册的是创建器，交给tensorRT管理
REGISTER_TENSORRT_PLUGIN(MySELUPluginCreator);

// 用于序列化插件的Helper function
template <typename T>
void writeToBuffer(char*& buffer, const T& val)
{
    *reinterpret_cast<T*>(buffer) = val;
    buffer += sizeof(T);
}

// 用于反序列化插件的Helper function
template <typename T>
T readFromBuffer(const char*& buffer)
{
    T val = *reinterpret_cast<const T*>(buffer);
    buffer += sizeof(T);
    return val;
}
// 定义插件类MYSELUPlugin
MySELUPlugin::MySELUPlugin(const std::string name, const std::string attr1, float attr3)
    : mLayerName(name),
    mattr1(attr1),
    mattr3(attr3)
{
    printf("==================== 编译阶段，attr1 = %s, attr3 = %f\n", attr1.c_str(), attr3);
}

MySELUPlugin::MySELUPlugin(const std::string name, const void* data, size_t length)
    : mLayerName(name)
{
    // Deserialize in the same order as serialization
    const char* d = static_cast<const char*>(data);
    const char* a = d;

    int nstr = readFromBuffer<int>(d);
    mattr1 = std::string(d, d + nstr);

    d += nstr;
    mattr3 = readFromBuffer<float>(d);
    assert(d == (a + length));

    printf("==================== 推理阶段，attr1 = %s, attr3 = %f\n", mattr1.c_str(), mattr3);
}

const char* MySELUPlugin::getPluginType() const noexcept
{
    return MYSELU_PLUGIN_NAME;
}

const char* MySELUPlugin::getPluginVersion() const noexcept
{
    return MYSELU_PLUGIN_VERSION;
}

int MySELUPlugin::getNbOutputs() const noexcept
{
    return 1;
}
// 获取该层的输出维度是多少
nvinfer1::DimsExprs MySELUPlugin::getOutputDimensions(int32_t outputIndex, const nvinfer1::DimsExprs* inputs, int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
{
    // Validate input arguments
    // assert(nbInputDims == 1);

    // MySELUping不改变输入尺寸，所以输出尺寸将与输入尺寸相同
    return *inputs;
}

int MySELUPlugin::initialize() noexcept
{
    return 0;
}
// 行性能测试
int MySELUPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc,
    const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
    void* output = outputs[0];
    size_t volume = 1;
    for (int i = 0; i < inputDesc->dims.nbDims; i++){
        volume *= inputDesc->dims.d[i];
    }
    mInputVolume = volume;

    myselu_inference(
        static_cast<const float*>(inputs[0]), 
        static_cast<float*>(output), 
        mInputVolume,
        stream
    );
    return 0;
}

size_t MySELUPlugin::getSerializationSize() const noexcept
{
    return sizeof(int) + mattr1.size() + sizeof(mattr3);
}
// 该层的参数序列化储存为trtmodel文件
void MySELUPlugin::serialize(void* buffer) const noexcept
{
    char* d = static_cast<char*>(buffer);
    const char* a = d;

    int nstr = mattr1.size();
    writeToBuffer(d, nstr);
    memcpy(d, mattr1.data(), nstr);

    d += nstr;
    writeToBuffer(d, mattr3);

    assert(d == a + getSerializationSize());
}
// 判断该插件所支持的数据格式和类型
bool MySELUPlugin::supportsFormatCombination(int32_t pos, const PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{   
    auto type = inOut[pos].type;
    auto format = inOut[pos].format;
    // 这个插件只支持普通的浮点数，以及NCHW输入格式
    if (type == DataType::kFLOAT && format == PluginFormat::kLINEAR)
        return true;
    else
        return false;
}

void MySELUPlugin::terminate() noexcept {}

void MySELUPlugin::destroy() noexcept
{
    // This gets called when the network containing plugin is destroyed
    delete this;
}
// 配置插件格式:告诉你目前这个层所采用的数据格式和类型
void MySELUPlugin::configurePlugin(const DynamicPluginTensorDesc* in, int32_t nbInputs,
    const DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept{
    // Validate input arguments

    auto type = in->desc.type;
    auto format = in->desc.format;
    assert(nbOutputs == 1);
    assert(type == DataType::kFLOAT);
    assert(format == PluginFormat::kLINEAR);
}
// 克隆插件
IPluginV2DynamicExt* MySELUPlugin::clone() const noexcept
{
    printf("===================克隆插件=================\n");
    auto plugin = new MySELUPlugin(mLayerName, mattr1, mattr3);
    plugin->setPluginNamespace(mNamespace.c_str());
    return plugin;
}

void MySELUPlugin::setPluginNamespace(const char* libNamespace) noexcept
{
    mNamespace = libNamespace;
}

const char* MySELUPlugin::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}
// 插件创建器
MySELUPluginCreator::MySELUPluginCreator()
{
    // 描述MySELUPlugin的必要PluginField参数
    mPluginAttributes.emplace_back(PluginField("attr1", nullptr, PluginFieldType::kCHAR, 0));
    mPluginAttributes.emplace_back(PluginField("attr3", nullptr, PluginFieldType::kFLOAT32, 1));

    // 收集PluginField的参数
    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

const char* MySELUPluginCreator::getPluginName() const noexcept
{
    return MYSELU_PLUGIN_NAME;
}

const char* MySELUPluginCreator::getPluginVersion() const noexcept
{
    return MYSELU_PLUGIN_VERSION;
}

const PluginFieldCollection* MySELUPluginCreator::getFieldNames() noexcept
{
    return &mFC;
}
// 创建plugin
IPluginV2* MySELUPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept
{
    std::string attr1;
    float attr3;
    const PluginField* fields = fc->fields;

    // Parse fields from PluginFieldCollection
    for (int i = 0; i < fc->nbFields; i++)
    {
        if (strcmp(fields[i].name, "attr1") == 0)
        {
            assert(fields[i].type == PluginFieldType::kCHAR);
            auto cp = static_cast<const char*>(fields[i].data);
            attr1 = std::string(cp, cp + fields[i].length);
        }
        else if (strcmp(fields[i].name, "attr3") == 0)
        {
            assert(fields[i].type == PluginFieldType::kFLOAT32);
            attr3 = *(static_cast<const float*>(fields[i].data));
        }
    }
    return new MySELUPlugin(name, attr1, attr3);
}
// 反序列化插件参数进行创建
IPluginV2* MySELUPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept
{
    // This object will be deleted when the network is destroyed, which will
    // call MySELUPlugin::destroy()
    return new MySELUPlugin(name, serialData, serialLength);
}

void MySELUPluginCreator::setPluginNamespace(const char* libNamespace) noexcept
{
    mNamespace = libNamespace;
}

const char* MySELUPluginCreator::getPluginNamespace() const noexcept
{
    return mNamespace.c_str();
}

通过 supportsFormatCombination 函数来告诉插件支持的数据格式和类型，enqueue 具体的 inference 就是在调用核函数，具体实现如下：

#include <cuda_runtime.h>
#include <math.h>

static __device__ float sigmoid(float x){
    return 1 / (1 + expf(-x));
}

static __global__ void myselu_kernel(const float* x, float* output, int n){

    int position = threadIdx.x + blockDim.x * blockIdx.x;
    if(position >= n) return;

    output[position] = x[position] * sigmoid(x[position]);
}

void myselu_inference(const float* x, float* output, int n, cudaStream_t stream){

    const int nthreads = 512;
    int block_size = n < nthreads ? n : nthreads;
    int grid_size = (n + block_size - 1) / block_size;
    myselu_kernel<<<grid_size, block_size, 0, stream>>>(x, output, n);
}

到此为止，我们大体走了一遍插件实现的流程，好难呀😭

2. 补充知识

关于插件你需要知道：(from 杜老师)

知识点

1. 插件主要是继承自 IPluginV2DynamicExt 后实现特定接口即可

实现插件要点

1. 导出 onnx 的时候，需要为 module 增加 symbolic 函数

• 参照这里：https://pytorch.org/docs/1.13/onnx.html#torch-autograd-functions
• g.op 对应的名称，需要与下面解析器的名称对应

2. src/onnx-tensorrt-release-8.0/builtin_op_importers.cpp:5094行，添加对插件 op 的解析

• DEFINE_BUILTIN_OP_IMPORTER(MYSELU)
• 注意解析时采用的名称要匹配上 src/myselu-plugin.cpp:15行

3. src/myselu-plugin.cpp:183行，创建 MySELUPluginCreator，插件创建器

• 实际注册时，注册的是创建器，交给 tensorRT 管理
• REGISTER_TENSORRT_PLUGIN(MySELUPluginCreator);
• src/myselu-plugin.cpp:23行

4. src/myselu-plugin.cpp:42行，定义插件类 MySELUPlugin

• Creator 创建器来实例化 MySELUPlugin 类

5. 正常使用该 onnx 即可

实现插件的两个阶段

1. 编译阶段

• 1. 通过 MySELUPluginCreator::createPlugin 创建 plugin
• 2. 期间会调用 MySELUPlugin::clone 克隆插件
• 3. 调用 MySELUPlugin::supportsFormatCombination 判断该插件所支持的数据格式和类型

• 在这里我们告诉引擎，本插件可以支持什么类型的推理

• 可以支持多种，例如 fp32、fp16、int8 等等

• 4. 调用 MySELUPlugin::getOutputDimension 获取该层的输出维度是多少
• 5. 调用 MySELUPlugin::enqueue 进行性能测试（不是一定会执行）

• 如果支持多种，则会在多种里面进行实际测试，选择一个性能最好的配置

• 6. 调用 MySELUPlugin::configurePlugin 配置插件格式

• 告诉你目前这个层所采用的数据格式和类型

• 7. 调用 MySELUPlugin::serialize 将该层的参数序列化储存为 trtmodel 文件

2. 推理阶段

• 1. 通过 MySELUPluginCreator::deserializePlugin 反序列化插件参数进行创建
• 2. 期间会调用 MySELUPlugin::clone 克隆插件
• 3. 调用 MySELUPlugin::configurePlugin 配置当前插件使用的数据类型和格式
• 4. 调用 MySELUPlugin::enqueue 进行推理

总结

本节课程学习了插件的编写，首先你要在 pytorch 中能够正确导出 onnx，自定义算子需要增加 symbolic 函数，然后你需要去 onnx-tensorrt 源码中的 buildin_op_importers 中添加对插件 op 的解析，接下来你要实现两个类，一个是插件创建类 Creator，另外一个是插件具体实现类，继承自 IPluginV2DynamicExt，其中具体推理代码还需要编写 CUDA 核函数完成，最后就是通过宏定义来注册插件。

感觉写插件的难度还是大，目前还没有完整的走过一遍流程🤣

而且一般来说一些简单的算子实现 tensorRT 应该都帮你搞定了，只有一些复杂的算子才需要写插件，流程虽然说是这么个流程，但是感觉复杂算子的插件难度上升得有点高，简单的都还没玩明白😂

非常硬核的技术活