TensorRT插件基础教程基于TensorRT8[2]

TensorRT插件零基础教程（2）（基于TensorRT8）

一、写在前面

引言

接着上一篇，本篇说明两个重要的类 nvinfer1::IPluginV2DynamicExt 和 IPluginCreator 的成员函数以及变量的定义，这两个重要的类实现之后，就基本完成了插件的编写。

读完本篇，你将获得

具体编写插件类的细节。

二、步入正题

1、定义派生的插件类

class GroupNormalizationPlugin final : public nvinfer1::IPluginV2DynamicExt

class GroupNormalizationPluginCreator : public IPluginCreator

final 表示这是最后一层继承，不允许再基于 GroupNormalizationPlugin 生成新的类。

2、两个插件内具体成员的定义

一般有两个构造函数，分别用于为插件内参数初始化、反序列化插件参数，由于在 C++ 里边，类有一个默认的构造函数，对于插件来说，显然这个默认的构造函数没有意义，所以手动使用 delete 保留字禁止其生成默认构造函数。

GroupNormalizationPlugin(float epsilon, const int nbGroups);

GroupNormalizationPlugin(const void* data, size_t length);

GroupNormalizationPlugin() = delete;

对于 GroupNormalizationPlugin 这个特定的插件，还有一些自定义的参数，在类内定义时用 private 修饰：

//通用内容
const char* mPluginNamespace;
std::string mNamespace;

//以下是自定义的，按需定义、使用
float mEpsilon;
int mNbGroups;
int mChannelVolume;

cudnnHandle_t _cudnn_handle; //这个插件用到了cuDNN库，需要定义这些内容
cudnnTensorDescriptor_t desc, bnDesc; // describes input and output
void* bnScale;
void* bnBias;

具体实现如下：

GroupNormalizationPlugin::GroupNormalizationPlugin(float epsilon, int nbGroups)
    : mEpsilon(epsilon),
      mNbGroups(nbGroups) //用初始化列表初始化参数
{
    //函数体内容按需写
    // Number of groups should be positive
    assert(nbGroups > 0);
}

//用于反序列化参数，即从引擎文件里读出保存的参数
GroupNormalizationPlugin::GroupNormalizationPlugin(const void* data, size_t length)
{
    // Deserialize in the same order as serialization
    deserialize_value(&data, &length, &mEpsilon); //这两个变量已经定义过了
    deserialize_value(&data, &length, &mNbGroups);
}

除了这两个构造函数以及一些自定义参数，其他的成员函数具有统一的格式。

//返回插件的名字，GROUP_NORM_NAME已经定义过了，很明确
const char* GroupNormalizationPlugin::getPluginType() const noexcept
{
    return GROUP_NORM_NAME;
}

//返回插件版本，已经定义过了，很明确，默认写1
const char* GroupNormalizationPlugin::getPluginVersion() const noexcept
{
    return GROUP_NORM_VERSION;
}

//获取该插件的输出数量
int GroupNormalizationPlugin::getNbOutputs() const noexcept
{
    return 1;
}

//从输入张量的维数得到计算输出张量维数的表达式，这个维度要用DimsExprs类表示
nvinfer1::DimsExprs GroupNormalizationPlugin::getOutputDimensions(
    int index, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
{
    // Input (from previous layer), scale and bias are the three inputs to the plugin.
    assert(nbInputs == 3); //检验传入参数的合法性
    assert(index == 0);  //检验传入参数的合法性
    nvinfer1::DimsExprs output(inputs[0]); //由于在本插件里，输出的维数与输入的维数相同，直接实例化一个DimsExprs，直接返回即可
    return output;
}

DimsExprs 的成员如下：

class DimsExprs
{
public:
    int32_t nbDims;                       //!< The number of dimensions.
    const IDimensionExpr* d[Dims::MAX_DIMS]; //!< The extent of each dimension.
};

//将插件对象附加到执行上下文，并授予插件访问某些上下文资源的权限
//此插件用到了cuDNN来实现插件的运算任务，该函数里的内容是为了配合cuDNN而存在
void GroupNormalizationPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
{
    _cudnn_handle = cudnnContext;
    cudnnCreateTensorDescriptor(&desc);
    cudnnCreateTensorDescriptor(&bnDesc);
}

//将插件对象从其执行上下文中分离出来
//此插件用到了cuDNN来实现插件的运算任务，该函数里的内容是为了配合cuDNN而存在
void GroupNormalizationPlugin::detachFromContext() noexcept
{
    cudnnDestroyTensorDescriptor(desc);
    cudnnDestroyTensorDescriptor(bnDesc);
}

下边的 enqueue 函数是重点，插件的实际运算函数从这里调用，因此用户需要仔细实现。此函数调用了 cuDNN 库里的函数，当然用户也可以自定义 CUDA 函数实现运算，留好接口，同样是在此处调用。

int GroupNormalizationPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
    const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace,
    cudaStream_t stream) noexcept
{
    // Get the input dimensions
    nvinfer1::Dims input_dims = inputDesc[0].dims;
    int batchSize = input_dims.d[0];
    int nbChannels = input_dims.d[1];

    // Calculate size of each group
    int groupSize = nbChannels / mNbGroups;

    mChannelVolume = std::accumulate(input_dims.d + 2, input_dims.d + inputDesc[0].dims.nbDims, 1, std::multiplies<int>());

    CHECK_CUDNN(cudnnSetTensor4dDescriptor(desc, // descriptor
        CUDNN_TENSOR_NCHW,                 // tensor format
        CUDNN_DATA_FLOAT,                  // type
        1,                                 // Batchsize
        batchSize * mNbGroups,             // Channels
        groupSize,                         // Height
        mChannelVolume                     // Width
        ));

    CHECK_CUDNN(cudnnDeriveBNTensorDescriptor(bnDesc, desc, CUDNN_BATCHNORM_SPATIAL));
    CHECK_CUDNN(cudnnSetStream(_cudnn_handle, stream));

    // Reshape the data according in the cudnnSetTensor4dDescriptor.
    float a = 1.F;
    float b = 0.F;
    CHECK_CUDNN(cudnnBatchNormalizationForwardTraining(_cudnn_handle, // handle
        CUDNN_BATCHNORM_SPATIAL,                                      // BatchNormMode_t, try also non persistent
        &a,                                                           //
        &b,                                                           //
        desc,                                                         // in/out descriptor
        inputs[0],                                                    // input
        desc,                                                         // in/out descriptor
        outputs[0],                                                   // output
        bnDesc,                                                       //
        bnScale,                                                      // 1
        bnBias,                                                       // 0
        0.0,                                                          // exponential average factor
        nullptr,                                                      // resultRunningMean
        nullptr,                                                      // resultRunningVar
        mEpsilon,                                                     //  eps
        nullptr,                                                      // resultSaveMean
        nullptr                                                       // resultSaveInvVar
        ));

    float* output = static_cast<float*>(outputs[0]);
    return scaleShiftChannelsInplace(output, batchSize, nbChannels, mChannelVolume, static_cast<const float*>(inputs[2]), static_cast<const float*>(inputs[1]), stream); //mBetaDev, mGammaDev,    
}

//获取序列化之后的参数总长度
size_t GroupNormalizationPlugin::getSerializationSize() const noexcept
{
    return sizeof(mNbGroups) + sizeof(mEpsilon);
}

//将参数序列化写入指定的buffer
void GroupNormalizationPlugin::serialize(void* buffer) const noexcept
{
    serialize_value(&buffer, mEpsilon);
    serialize_value(&buffer, mNbGroups);
}

//如果插件支持 pos 索引的输入/输出的格式和数据类型，则返回 true
//inOut指的是插件的输入张量或者输出张量的描述，通常判断提供过来的输入或者输出format以及type受不受插件支持
bool GroupNormalizationPlugin::supportsFormatCombination(
    int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) noexcept
{
    assert(inOut && pos < (nbInputs + nbOutputs));
    return ((inOut[pos].type == nvinfer1::DataType::kFLOAT) && inOut[pos].format == nvinfer1::PluginFormat::kLINEAR
        && inOut[pos].type == inOut[0].type); 
}

//释放插件层初始化过程中获得的资源，由于cpu并不主动释放gpu端的资源，所以要手动释放
void GroupNormalizationPlugin::terminate() noexcept
{
    cudaFree(bnScale);
    cudaFree(bnBias);
}

//销毁插件
void GroupNormalizationPlugin::destroy() noexcept
{
    // This gets called when the network containing plugin is destroyed
    delete this;
}

//克隆插件对象。这也会复制内部插件参数，并返回一个带有这些参数的新插件对象
IPluginV2DynamicExt* GroupNormalizationPlugin::clone() const noexcept
{
    auto* plugin = new GroupNormalizationPlugin(mEpsilon, mNbGroups);
    plugin->setPluginNamespace(mPluginNamespace); //必写
    return plugin;
}

//创建插件的准备工作，例如分配空间，数据传输等。
//configurePlugin ()可以在构建和执行阶段中多次调用。构建阶段发生在 initialize ()被调用之前，并且只发生在 IBuilder 创建引擎的过程中。执行阶段发生在 initialize ()被调用之后，并且在 IBuilder 创建引擎和 IExecutionContext 执行引擎期间发生。
void GroupNormalizationPlugin::configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs,
    const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) noexcept
{

    for (int i = 0; i < nbInputs; i++)
    {
      for (int j = 0; j < in[0].desc.dims.nbDims; j++)
      {
        // Do not support dynamic dimensions
        assert(in[0].desc.dims.d[j] != -1);
      }
    }

    int batchSize = in[0].desc.dims.d[0];
    int nbChannels = in[0].desc.dims.d[1];

    // Allocate device memory and initialize scale and bias values
    cudaMalloc(&bnScale, batchSize * nbChannels * sizeof(float));
    cudaMalloc(&bnBias, batchSize * nbChannels * sizeof(float));

    // allot ones and zeros to bn parameters
    std::vector<float> ones(nbChannels, 1.F);
    cudaMemcpy(bnScale, ones.data(), nbChannels * sizeof(float), cudaMemcpyHostToDevice);

    std::vector<float> zeroes(nbChannels, 0.F);
    cudaMemcpy(bnBias, zeroes.data(), nbChannels * sizeof(float), cudaMemcpyHostToDevice);
}

//在请求的索引处返回插件输出的 DataType
//默认行为应该是返回第一个输入的类型，如果层没有输入，则返回 DataType: : : kFLOAT
nvinfer1::DataType GroupNormalizationPlugin::getOutputDataType(
    int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept
{
    assert(inputTypes && nbInputs > 0 && index == 0);
    return inputTypes[0];
}

//获取插件/也可以说为层所需的工作空间
//TODO 这里所说的工作空间暂不明确到底是什么，configurePlugin函数已经分配好了空间
size_t GroupNormalizationPlugin::getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs,
    const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const noexcept
{
    return 0;
}

//设置此插件对象所属的名称空间。理想情况下，来自同一个插件库的所有插件对象应该具有相同的名称空间。
void GroupNormalizationPlugin::setPluginNamespace(const char* libNamespace) noexcept
{
    mPluginNamespace = libNamespace;
}

//获取插件的命名空间
const char* GroupNormalizationPlugin::getPluginNamespace() const noexcept
{
    return mPluginNamespace;
}

上述所有函数虽然比较繁多，但是需要认真琢磨的也为数不多，下面的 GroupNormalizationPluginCreator 类也是如此：

//用于初始化参数的构造函数
GroupNormalizationPluginCreator::GroupNormalizationPluginCreator()
{
    mPluginAttributes.clear();
    mPluginAttributes.emplace_back(PluginField("eps", nullptr, PluginFieldType::kFLOAT32, 1));
    mPluginAttributes.emplace_back(PluginField("num_groups", nullptr, PluginFieldType::kINT32, 1));

    mFC.nbFields = mPluginAttributes.size();
    mFC.fields = mPluginAttributes.data();
}

//以下函数表意明确，结构简单，照猫画虎即可，但不能省略
const char* GroupNormalizationPluginCreator::getPluginName() const noexcept
{
    return GROUP_NORM_NAME;
}

const char* GroupNormalizationPluginCreator::getPluginVersion() const noexcept
{
    return GROUP_NORM_VERSION;
}

const PluginFieldCollection* GroupNormalizationPluginCreator::getFieldNames() noexcept
{
    return &mFC;
}

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

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

//用于创建插件的函数，从PluginFieldCollection中解析出插件的fields(即为保存有插件信息的内存空间）
IPluginV2DynamicExt* GroupNormalizationPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept
{
    // Set default values
    int nbGroups{1};
    float epsilon{0.00001F};
    for (int i = 0; i < fc->nbFields; i++)   //遍历解析
    {
        std::string field_name(fc->fields[i].name);
        if (field_name.compare("eps") == 0)  //fields中的name是字符串
        {
            epsilon = *static_cast<const float*>(fc->fields[i].data);//data是void* 类型，需要转为数值
        }
        if (field_name.compare("num_groups") == 0)
        {
            nbGroups = *static_cast<const int*>(fc->fields[i].data);
        }
    }

    GroupNormalizationPlugin* plugin = new GroupNormalizationPlugin(epsilon, nbGroups);
    plugin->setPluginNamespace(mNamespace.c_str());

    return plugin;
}

fields结构体如下：

struct PluginFieldCollection
{
    //! Number of PluginField entries.
    int32_t nbFields;
    //! Pointer to PluginField entries.
    PluginField const* fields;
};
class PluginField
{
public:
    AsciiChar const* name;
    void const* data;
    PluginFieldType type;
    int32_t length;
    PluginField(AsciiChar const* const name_ = nullptr, void const* const data_ = nullptr,
        PluginFieldType const type_ = PluginFieldType::kUNKNOWN, int32_t const length_ = 0) noexcept
        : name(name_)
        , data(data_)
        , type(type_)
        , length(length_)
    {
    }
};

//反序列化，data就是需要解析出来的数据，length是总长度
IPluginV2DynamicExt* GroupNormalizationPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept
{
    GroupNormalizationPlugin* plugin = new GroupNormalizationPlugin(serialData, serialLength);
    plugin->setPluginNamespace(mNamespace.c_str());

    return plugin;
}