五. TensorRT API的基本使用-custom-trt-plugin

前言

自动驾驶之心推出的 《CUDA与TensorRT部署实战课程》，链接。记录下个人学习笔记，仅供自己参考

本次课程我们来学习课程第五章—TensorRT API 的基本使用，一起来学习自定义插件的编写

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

0. 简述

本小节目标：学习 TensorRT 自定义插件的编写

今天我们来讲第五章节第七小节—5.7-custom-basic-trt-plugin 这个案例

下面我们开始本次课程的学习🤗

1. 案例运行

在正式开始课程之前，博主先带大家跑通 5.7-custom-basic-trt-plugin 这个小节的案例🤗
源代码获取地址：https://github.com/kalfazed/tensorrt_starter

首先大家需要把 tensorrt_starter 这个项目给 clone 下来，指令如下：

git clone https://github.com/kalfazed/tensorrt_starter.git

也可手动点击下载，点击右上角的 Code 按键，将代码下载下来。至此整个项目就已经准备好了。也可以点击 here 下载博主准备好的源代码（注意代码下载于 2024/7/14 日，若有改动请参考最新）

整个项目后续需要使用的软件主要有 CUDA、cuDNN、TensorRT、OpenCV，大家可以参考 Ubuntu20.04软件安装大全 进行相应软件的安装，博主这里不再赘述

假设你的项目、环境准备完成，下面我们来一起运行 5.7 小节案例代码

开始之前我们需要创建几个文件夹，在 tensorrt_starter/chapter5-tensorrt-api-basics/5.7-custom-basic-trt-plugin 小节中创建一个 models 文件夹，接着在 models 文件夹下创建 onnx 和 engine 文件夹，总共三个文件夹需要创建

创建完后 5.7 小节整个目录结构如下：

接着我们需要执行 python 文件创建一个 ONNX 模型，先进入到 5.7 小节中：

cd tensorrt_starter/chapter5-tensorrt-api-basics/5.7-custom-basic-trt-plugin

执行如下指令：

python src/python/plugin_customScalar.py

Note：大家需要准备一个虚拟环境，安装好 torch、onnx、onnxsim 等第三方库

输出如下：

生成好的 onnx 模型文件保存在 models/onnx 文件夹下，大家可以查看

接着我们需要利用 ONNX 生成对应的 engine，在此之前我们需要修改下整体的 Makefile.config，指定一些库的路径：

# tensorrt_starter/config/Makefile.config
# CUDA_VER                    :=  11
CUDA_VER                    :=  11.6
    
# opencv和TensorRT的安装目录
OPENCV_INSTALL_DIR          :=  /usr/local/include/opencv4
# TENSORRT_INSTALL_DIR        :=  /mnt/packages/TensorRT-8.4.1.5
TENSORRT_INSTALL_DIR        :=  /home/jarvis/lean/TensorRT-8.6.1.6

Note：大家查看自己的 CUDA 是多少版本，修改为对应版本即可，另外 OpenCV 和 TensorRT 修改为你自己安装的路径即可

接着我们就可以来执行编译，指令如下：

make -j64

输出如下：

接着执行：

./trt-infer

输出如下：

我们这里通过加载带有自定义节点 scalar 的 onnx 模型生成对应的 engine，并加载实现自定义节点 scalar 的插件完成推理，可以看到和 python 推理结果保持一致

如果大家能够看到上述输出结果，那就说明本小节案例已经跑通，下面我们就来看看具体的代码实现

2. 代码分析

2.1 plugin_customScalar.py

我们先来看下对应的 python 脚本文件都做了些啥：

import torch
import torch.onnx
import torch.nn as nn
import onnxruntime
import onnx
import onnxsim
import os
from collections import OrderedDict

class CustomScalarImpl(torch.autograd.Function):
    @staticmethod
    def symbolic(g, x, r, s):
        return g.op("custom::customScalar", x, scalar_f=r, scale_f=s)

    @staticmethod
    def forward(ctx, x, r, s):
        return (x + r) * s

class CustomScalar(nn.Module):
    def __init__(self, r, s):
        super().__init__()
        self.scalar = r
        self.scale  = s

    def forward(self, x):
        return CustomScalarImpl.apply(x, self.scalar, self.scale)


class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.conv   = nn.Conv2d(1, 3, (3, 3), padding=1)
        self.act    = CustomScalar(1, 10)
    
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, mean=0., std=1.)
            if isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
                nn.init.constant_(m.weight, 1.05)
                nn.init.constant_(m.bias, 0.05)

    def forward(self, x):
        x = self.conv(x)
        x = self.act(x)
        return x

def setup_seed(seed):
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True

def export_norm_onnx(input, model):
    current_path = os.path.dirname(__file__)
    file = current_path + "/../../models/onnx/sample_customScalar.onnx"
    torch.onnx.export(
        model         = model, 
        args          = (input,),
        f             = file,
        input_names   = ["input0"],
        output_names  = ["output0"],
        opset_version = 15)
    print("Finished normal onnx export")

    # check the exported onnx model
    model_onnx = onnx.load(file)
    onnx.checker.check_model(model_onnx)

    # use onnx-simplifier to simplify the onnx
    print(f"Simplifying with onnx-simplifier {onnxsim.__version__}...")
    model_onnx, check = onnxsim.simplify(model_onnx)
    assert check, "assert check failed"
    onnx.save(model_onnx, file)

def eval(input, model):
    output = model(input)
    print("------from infer------")
    print(input)
    print("\n")
    print(output)

if __name__ == "__main__":
    setup_seed(1)
    input = torch.tensor([[[
        [0.7576, 0.2793, 0.4031, 0.7347, 0.0293],
        [0.7999, 0.3971, 0.7544, 0.5695, 0.4388],
        [0.6387, 0.5247, 0.6826, 0.3051, 0.4635],
        [0.4550, 0.5725, 0.4980, 0.9371, 0.6556],
        [0.3138, 0.1980, 0.4162, 0.2843, 0.3398]]]])

    model = Model()
    model.eval() 
    
    # 计算
    eval(input, model)

    # 导出onnx
    export_norm_onnx(input, model);

该脚本定义了一个自定义的 PyTorch 插件，并展示了如何将其导出 ONNX，以下是对核心内容的简单分析：

1. 自定义 Scalar 插件

class CustomScalarImpl(torch.autograd.Function):
    @staticmethod
    def symbolic(g, x, r, s):
        return g.op("custom::customScalar", x, scalar_f=r, scale_f=s)

    @staticmethod
    def forward(ctx, x, r, s):
        return (x + r) * s

• CustomScalarImpl 继承自 torch.autograd.Function，定义了一个自定义的 PyTorch 操作
• symbolic 方法：定义了如何将这个自定义操作转换为 ONNX 图中的节点，它创建了一个名为 custom::customScalar 的 ONNX 操作，带有 scalar_f 和 scale_f 两个参数
• forward 方法：实现了自定义操作的前向计算逻辑，即对输入 x 加上 r 然后乘以 s

2. 自定义 Scalar 模块

class CustomScalar(nn.Module):
    def __init__(self, r, s):
        super().__init__()
        self.scalar = r
        self.scale  = s

    def forward(self, x):
        return CustomScalarImpl.apply(x, self.scalar, self.scale)

• CustomScalar 是一个 PyTorch 模块，使用 CustomScalarImpl 来实现自定义的前向计算
• __init__ 方法：初始化自定义标量 r 和缩放因子 s
• forward 方法：使用 CustomScalarImpl.apply 调用自定义的前向计算逻辑

这里博主其实有个困惑，那就是 CustomScalar 模块的 forward 过程为什么不直接调用 CustomScalarImpl 的 forward 而是去调用 CustomScalarImpl 的 apply 方法呢？🤔

apply 是 torch.autograd.Function 提供的一个特殊方法，用于执行自定义操作，它会调用 forward 方法来进行前向计算，使用 apply 方法主要是基于以下几点原因：

• 自动计算图构建：apply 方法会自动构建计算图，并将操作添加到图中，以便在反向传播时能够计算梯度
• 内部处理：apply 方法处理一些内部细节，如将参数传递给 forward 方法，并确保计算图的正确性
• 兼容性：使用 apply 方法可以确保自定义操作与 PyTorch 的自动微分功能兼容

接着我们来看下导出的 ONNX 的结构：

可以看到导出的 ONNX 中有我们自定义的 customScalar 节点，该节点的 op_type 就是 customScalar，attributes 属性包括 scalar 和 scale

这里有个点需要大家注意，那就是我们通过 g.op 创建 op 算子时，参数的传递例如 scalar 和 scale 我们都加上了 _f 代表这两个参数的类型是 float，如果是 int 类型那就加上 _i 就行

下面我们来看 C++ 中是如何实现 customScalar 的

2.2 main.cpp

我们先从 main.cpp 看起：

#include <iostream>
#include <memory>

#include "utils.hpp"
#include "model.hpp"

using namespace std;

int main(int argc, char const *argv[])
{
    Model model("models/onnx/sample_customScalar.onnx", Model::precision::FP32);
    // Model model("models/onnx/sample_customLeakyReLU.onnx", Model::precision::FP16);
    if(!model.build()){
        LOGE("fail in building model");
        return 0;
    }
    if(!model.infer()){
        LOGE("fail in infering model");
        return 0;
    }
    return 0;
}

main 函数中提供了 customScalar 和 customLeakyReLU 两个案例，Model 构造函数传入 ONNX 和指定精度，接着通过 build 接口构建 engine，infer 接口推理

2.3 model.cpp

我们先来看下 Model 类构造函数：

Model::Model(string path, precision prec){
    if (getFileType(path) == ".onnx")
        mOnnxPath = path;
    else if (getFileType(path) == ".weights")
        mWtsPath = path;
    else 
        LOGE("ERROR: %s, wrong weight or model type selected. Program terminated", getFileType(path).c_str());

    if (prec == precision::FP16) {
        mPrecision = nvinfer1::DataType::kHALF;
    } else if (prec == precision::INT8) {
        mPrecision = nvinfer1::DataType::kINT8;
    } else {
        mPrecision = nvinfer1::DataType::kFLOAT;
    }

    mEnginePath = getEnginePath(path, prec);
}

和之前的构造函数一样，首先它根据文件扩展名判断模型文件类型，分别设置 ONNX 模型路径和权重文件路径，然后根据指定的精度参数来设置不同的模型精度，并生成对应的 engine 路径用于后续模型构建和加载

接着看下 build 接口：

bool Model::build_from_onnx(){
    if (fileExists(mEnginePath)){
        LOG("%s has been generated!", mEnginePath.c_str());
        return true;
    } else {
        LOG("%s not found. Building engine...", mEnginePath.c_str());
    }
    Logger logger;
    auto builder       = make_unique<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(logger));
    auto network       = make_unique<nvinfer1::INetworkDefinition>(builder->createNetworkV2(1));
    auto config        = make_unique<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
    auto parser        = make_unique<nvonnxparser::IParser>(nvonnxparser::createParser(*network, logger));

    config->setMaxWorkspaceSize(1<<28);
    config->setProfilingVerbosity(nvinfer1::ProfilingVerbosity::kDETAILED);

    if (!parser->parseFromFile(mOnnxPath.c_str(), 1)){
        LOGE("ERROR: failed to %s", mOnnxPath.c_str());
        return false;
    }

    if (builder->platformHasFastFp16() && mPrecision == nvinfer1::DataType::kHALF) {
        config->setFlag(nvinfer1::BuilderFlag::kFP16);
        config->setFlag(nvinfer1::BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
    } else if (builder->platformHasFastInt8() && mPrecision == nvinfer1::DataType::kINT8) {
        config->setFlag(nvinfer1::BuilderFlag::kINT8);
        config->setFlag(nvinfer1::BuilderFlag::kPREFER_PRECISION_CONSTRAINTS);
    }

    auto engine        = make_unique<nvinfer1::ICudaEngine>(builder->buildEngineWithConfig(*network, *config));
    auto plan          = builder->buildSerializedNetwork(*network, *config);
    auto runtime       = make_unique<nvinfer1::IRuntime>(nvinfer1::createInferRuntime(logger));

    auto f = fopen(mEnginePath.c_str(), "wb");
    fwrite(plan->data(), 1, plan->size(), f);
    fclose(f);

    // mEngine            = shared_ptr<nvinfer1::ICudaEngine>(runtime->deserializeCudaEngine(plan->data(), plan->size()));
    // mInputDims         = network->getInput(0)->getDimensions();
    // mOutputDims        = network->getOutput(0)->getDimensions();

    // 把优化前和优化后的各个层的信息打印出来
    // LOG("Before TensorRT optimization");
    // print_network(*network, false);
    // LOG("");
    // LOG("After TensorRT optimization");
    // print_network(*network, true);

    // LOG("Finished building engine");
    return true;
};

build 接口和之前的差别不大，插件的具体体现是在 parser 上，我们通过 parseFromFile 将 ONNX 各个节点信息 parse 到 network 中，这个时候如果我们不做一些处理，tensorRT 是没有办法识别我们自定义的 customScalar 算子的

2.4 custom-scalar-plugin

这里我们先来看看自定义插件 scalar 的头文件：

下面是自定义插件 Scalar 的头文件，请你帮我详细分析下这个头文件都做了哪些内容，同时帮我总结下在 C++ 中要实现一个自定义插件我们应该怎么做，具体都需要包含哪些内容：

#ifndef __CUSTOM_SCARLAR_PLUGIN_HPP
#define __CUSTOM_SCARLAR_PLUGIN_HPP

#include "NvInferRuntime.h"
#include "NvInferRuntimeCommon.h"
#include <NvInfer.h>
#include <string>
#include <vector>

using namespace nvinfer1;

namespace custom 
{
static const char* PLUGIN_NAME {"customScalar"};
static const char* PLUGIN_VERSION {"1"};


/* 
* 在这里面需要创建两个类, 一个是普通的Plugin类, 一个是PluginCreator类
*  - Plugin类是插件类，用来写插件的具体实现
*  - PluginCreator类是插件工厂类，用来根据需求创建插件。调用插件是从这里走的
*/

class CustomScalarPlugin : public IPluginV2DynamicExt {
public:
    /*
     * 我们在编译的过程中会有大概有三次创建插件实例的过程
     * 1. parse阶段: 第一次读取onnx来parse这个插件。会读取参数信息并转换为TensorRT格式
     * 2. clone阶段: parse完了以后，TensorRT为了去优化这个插件会复制很多副本出来来进行很多优化测试。也可以在推理的时候供不同的context创建插件的时候使用
     * 3. deseriaze阶段: 将序列化好的Plugin进行反序列化的时候也需要创建插件的实例
    */
    CustomScalarPlugin() = delete; //默认构造函数，一般直接delete
    CustomScalarPlugin(const std::string &name, float scalar, float scale);  //parse 时候用的构造函数
    CustomScalarPlugin(const std::string &name, const void* buffer, size_t length); //clone 以及反序列化的时候用的构造函数

    ~CustomScalarPlugin();

    /* 有关获取plugin信息的方法 */
    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; //查看pos位置的索引是否支持指定的DataType以及TensorFormat
    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 scalar;
        float scale;
    } mParams; // 当这个插件op需要有参数的时候，把这些参数定义为成员变量，可以单独拿出来定义，也可以像这样定义成一个结构体
};

class CustomScalarPluginCreator : public IPluginCreator {
public:
    CustomScalarPluginCreator();  //初始化mFC以及mAttrs
    ~CustomScalarPluginCreator();

    const char*                     getPluginName() const noexcept override;
    const char*                     getPluginVersion() const noexcept override;
    const PluginFieldCollection*    getFieldNames() noexcept override;
    const char*                     getPluginNamespace() const noexcept override;
    IPluginV2*                      createPlugin(const char* name, const PluginFieldCollection* fc) noexcept override;  //通过包含参数的mFC来创建Plugin。调用上面的Plugin的构造函数
    IPluginV2*                      deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept override;
    void                            setPluginNamespace(const char* pluginNamespace) noexcept override;
      
private:
    static PluginFieldCollection    mFC;           //接受plugionFields传进来的权重和参数，并将信息传递给Plugin，内部通过createPlugin来创建带参数的plugin
    static std::vector<PluginField> mAttrs;        //用来保存这个插件op所需要的权重和参数, 从onnx中获取, 同样在parse的时候使用
    std::string                     mNamespace;
    
};

} // namespace custom

#endif __CUSTOM_SCARLAR_PLUGIN_HPP

这个头文件定义了一个名为 CustomScalarPlugin 的自定义插件，以及与之相关的工厂类 CustomScalarPluginCreator，这些类是为了在 TensorRT 中实现特定功能而创建的

写插件时我们要遵循几个原则，首先是定义命名空间，这个其实是跟 python 中对齐的，例如这里的 namespace custom 和 python 中的 custom::customScalar 对应。接着插件的名字和版本也需要定义，名字也是和 python 中对齐的

接着我们在命名空间下创建了两个类，分别是 CustomScalarPlugin 和 CustomScalarPluginCreator，它们有着各自的功能：(from ChatGPT)

CustomScalarPlugin 类功能

• 1. 插件执行：CustomScalarPlugin 是核心类，负责实现具体的运算逻辑。这个类直接处理输入和输出张量，执行所需的计算
• 2. 序列化支持：支持插件的序列化和反序列化功能，这是保存和加载训练好的网络模型所必需的
• 3. 资源管理：管理与插件相关的资源，例如内存和计算资源

CustomScalarPluginCreator 类功能

• 1. 插件工厂：这个类充当插件的工厂，负责根据定义的参数创建插件的实例。它读取网络定义并根据这些信息创建 CustomScalarPlugin 实例
• 2. 插件注册：它还负责将插件注册到 TensorRT 插件注册表中，以便 TensorRT 在解析模型时能够识别并正确地创建插件实例
• 3. 参数处理：管理和解析用于插件构造的参数，这包括从网络定义中读取参数，并将它们传递给 CustomScalarPlugin 构造函数

简单来说 Plugin 类是插件类，用来写插件的具体实现，PluginCreator 类是插件工厂类，用来根据需求创建插件，调用插件是从这里走的

我们来看 CustomScalarPluginCreator 类，它在构造函数中需要初始化 mFC 和 mAttrs：

• mFC
  • type 是 PluginFieldCollection
  • 它接收 pluginFields 传进来的权重和参数，并将信息传递给 Plugin，内部通过 createPlugin 来创建带有参数的 plugin
• mAttrs
  • type 是 vector<PluginField>
  • 它用来保存这个插件 op 所需要的权重和参数，从 onnx 中获取同样在 parse 的时候使用

其构造函数具体实现如下：

CustomScalarPluginCreator::CustomScalarPluginCreator()
{
    /* 
     * 每个插件的Creator构造函数需要定制，主要就是获取参数以及传递参数
     * 初始化creator中的PluginField以及PluginFieldCollection
     * - PluginField::            负责获取onnx中的参数
     * - PluginFieldCollection：  负责将onnx中的参数传递给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();
}

其它函数的实现如下：

CustomScalarPluginCreator::~CustomScalarPluginCreator()
{
    /* 一般不需要做任何使用，所有插件实现都差不多 */
}

const char* CustomScalarPluginCreator::getPluginName() const noexcept
{
    /* 所有插件实现都差不多 */
    return PLUGIN_NAME;
}

const char* CustomScalarPluginCreator::getPluginVersion() const noexcept 
{
    /* 所有插件实现都差不多 */
    return PLUGIN_VERSION;
}

const char* CustomScalarPluginCreator::getPluginNamespace() const noexcept
{
    /* 所有插件实现都差不多 */
    return mNamespace.c_str();
}

IPluginV2* CustomScalarPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept 
{
    /*
     * 通过Creator创建一个Plugin的实现，这个时候会通过mFC中取出需要的参数, 并实例化一个Plugin
     * 这个案例中，参数有scalar和scale两个参数。从fc中取出来对应的数据来初始化这个plugin
    */
    float scalar = 0;
    float scale  = 0;
    std::map<std::string, float*> paramMap = {{"scalar", &scalar}, {"scale", &scale}};

    for (int i = 0; i < fc->nbFields; i++) {
        if (paramMap.find(fc->fields[i].name) != paramMap.end()){
            *paramMap[fc->fields[i].name] = *reinterpret_cast<const float*>(fc->fields[i].data);
        }
    }
    return new CustomScalarPlugin(name, scalar, scale);
}

IPluginV2* CustomScalarPluginCreator::deserializePlugin(const char* name, const void* serialData, size_t serialLength) noexcept
{
    /* 反序列化插件其实就是实例化一个插件，所有插件实现都差不多 */
    return new CustomScalarPlugin(name, serialData, serialLength);
}

void CustomScalarPluginCreator::setPluginNamespace(const char* pluginNamespace) noexcept
{
    /* 所有插件实现都差不多 */
    mNamespace = pluginNamespace;
    return;
}

const PluginFieldCollection* CustomScalarPluginCreator::getFieldNames() noexcept
{
    /* 所有插件实现都差不多 */
    return &mFC;
}

从上面函数的实现可以看到很多函数的实现其实都比较简单，我们重点来看下 createPlugin 函数：

IPluginV2* CustomScalarPluginCreator::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept 
{
    /*
     * 通过Creator创建一个Plugin的实现，这个时候会通过mFC中取出需要的参数, 并实例化一个Plugin
     * 这个案例中，参数有scalar和scale两个参数。从fc中取出来对应的数据来初始化这个plugin
    */
    float scalar = 0;
    float scale  = 0;
    std::map<std::string, float*> paramMap = {{"scalar", &scalar}, {"scale", &scale}};

    for (int i = 0; i < fc->nbFields; i++) {
        if (paramMap.find(fc->fields[i].name) != paramMap.end()){
            *paramMap[fc->fields[i].name] = *reinterpret_cast<const float*>(fc->fields[i].data);
        }
    }
    return new CustomScalarPlugin(name, scalar, scale);
}

该函数根据传入的名称和参数集合 PluginFieldCollection 创建一个新的 CustomScalarPlugin 实例。PluginFieldCollection 包含了插件所需的所有参数，这些参数在模型解析阶段被从模型文件中提取：(from ChatGPT)

• 1. 局部变量定义
  • scalar 和 scale 是插件需要的两个参数，分别初始化为0
• 2. 参数映射
  • 使用 std::map<std::string, float*> paramMap 来建立参数名到这些变量地址的映射，这使得通过参数名访问和设置这些参数变得容易
• 3. 遍历并设置参数
  • 通过遍历 fc（即 PluginFieldCollection）中的所有字段来检查哪些参数被提供
  • 对于每个字段，如果它的名称存在于 paramMap 中，则将该字段的值（通过 reinterpret_cast 转换为 float*）赋给相应的变量
• 4. 创建并返回插件实例
  • 使用解析出的参数 scalar 和 scale 调用 CustomScalarPlugin 的构造函数来创建一个新的插件实例

我们分析完 CustomScalarPluginCreator 我们再来看 CustomScalarPlugin，CustomScalarPlugin 有三个不同的构造函数，我们在编译的过程中大概有三次创建插件实例的过程：

• 1. parser 阶段：第一次读取 onnx 来 parse 这个插件会读取参数信息并转换为 TensorRT 的格式
• 2. clone 阶段：parse 之后，TensorRT 为了去优化插件会复制很多副本来进行测试，也可以在推理的时候提供不同的 context 创建插件的时候使用
• 3. deserialize 阶段：将序列化好的 Plugin 进行反序列化的时候也需要创建插件的实例

其构造函数具体实现如下：

CustomScalarPlugin::CustomScalarPlugin(const std::string &name, float scalar, float scale):
    mName(name)
{
    mParams.scalar = scalar;
    mParams.scale = scale;
}

CustomScalarPlugin::CustomScalarPlugin(const std::string &name, const void* buffer, size_t length):
    mName(name)
{
    memcpy(&mParams, buffer, sizeof(mParams));
}

其它函数的实现如下：

CustomScalarPlugin::~CustomScalarPlugin()
{
    /* 这里的析构函数不需要做任何事情，生命周期结束的时候会自动调用terminate和destroy */
    return;
}

const char* CustomScalarPlugin::getPluginType() const noexcept
{
    /* 一般来说所有插件的实现差不多一致 */
    return PLUGIN_NAME;
}

const char* CustomScalarPlugin::getPluginVersion() const noexcept
{
    /* 一般来说所有插件的实现差不多一致 */
    return PLUGIN_VERSION;
}

int32_t CustomScalarPlugin::getNbOutputs() const noexcept
{
    /* 一般来说所有插件的实现差不多一致 */
    return 1;
}

size_t CustomScalarPlugin::getSerializationSize() const noexcept
{
    /* 如果把所有的参数给放在mParams中的话, 一般来说所有插件的实现差不多一致 */
    return sizeof(mParams);
}

const char* CustomScalarPlugin::getPluginNamespace() const noexcept
{
    /* 一般来说所有插件的实现差不多一致 */
    return mNamespace.c_str();
}

DataType CustomScalarPlugin::getOutputDataType(int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
{
    /* 一般来说所有插件的实现差不多一致 */
    return inputTypes[0];
}

DimsExprs CustomScalarPlugin::getOutputDimensions(int32_t outputIndex, const DimsExprs* inputs, int32_t nbInputs, IExprBuilder &exprBuilder) noexcept
{
    /* 一般来说所有插件的实现差不多一致 */
    return inputs[0];
}

size_t CustomScalarPlugin::getWorkspaceSize(const PluginTensorDesc *inputs, int32_t nbInputs, const PluginTensorDesc *outputs, int32_t nbOutputs) const noexcept
{
    /* 一般来说会使用builder创建时用的workspaceSize所以这里一般什么都不做 */
    return 0;
}

int32_t CustomScalarPlugin::initialize() noexcept
{
    /* 这个一般会根据情况而定，建议每个插件都有一个自己的实现 */
    return 0;
}

void CustomScalarPlugin::terminate() noexcept 
{
    /* 
     * 这个是析构函数调用的函数。一般和initialize配对的使用
     * initialize分配多少内存，这里就释放多少内存
    */
    return;
}

void CustomScalarPlugin::serialize(void *buffer) const noexcept
{
    /* 序列化也根据情况而定，每个插件自己定制 */
    memcpy(buffer, &mParams, sizeof(mParams));
    return;

}

void CustomScalarPlugin::destroy() noexcept
{
    /* 一般来说所有插件的实现差不多一致 */
    delete this;
    return;
}

int32_t CustomScalarPlugin::enqueue(
    const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, 
    const void* const* inputs, void* const* outputs, 
    void* workspace, cudaStream_t stream) noexcept
{
    /*
     * Plugin的核心的地方。每个插件都有一个自己的定制方案
     * Plugin直接调用kernel的地方
    */
    int nElements = 1;
    for (int i = 0; i < inputDesc[0].dims.nbDims; i++){
        nElements *= inputDesc[0].dims.d[i];
    }

    customScalarImpl(
            static_cast<const float*>(inputs[0]),
            static_cast<float*>(outputs[0]), 
            mParams.scalar, 
            mParams.scale,
            nElements,
            stream);

    return 0;
}

IPluginV2DynamicExt* CustomScalarPlugin::clone() const noexcept
{
    /* 克隆一个Plugin对象，所有的插件的实现都差不多*/
    auto p = new CustomScalarPlugin(mName, &mParams, sizeof(mParams));
    p->setPluginNamespace(mNamespace.c_str());
    return p;
}

bool CustomScalarPlugin::supportsFormatCombination(int32_t pos, const PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
    /* 
     * 设置这个Plugin支持的Datatype以及TensorFormat, 每个插件都有自己的定制
     * 作为案例展示，这个customScalar插件只支持FP32，如果需要扩展到FP16以及INT8，需要在这里设置
    */
    
    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;
}

void CustomScalarPlugin::configurePlugin(const DynamicPluginTensorDesc* in, int32_t nbInputs, const DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept
{
    /* 一般不需要做任何使用，所有插件实现都差不多 */
    return;
}
void CustomScalarPlugin::setPluginNamespace(const char* pluginNamespace) noexcept
{
    /* 所有插件的实现都差不多 */
    mNamespace = pluginNamespace;
    return;
}
void CustomScalarPlugin::attachToContext(cudnnContext* contextCudnn, cublasContext* contextCublas, IGpuAllocator *gpuAllocator) noexcept 
{
    /* 一般不需要做任何使用，所有插件实现都差不多 */
    return;
}
void CustomScalarPlugin::detachFromContext() noexcept 
{
    /* 一般不需要做任何使用，所有插件实现都差不多 */
    return;
}

从上面函数的实现可以看到很多函数的实现其实都比较简单，我们重点来看下 enqueue 函数：

int32_t CustomScalarPlugin::enqueue(
    const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc, 
    const void* const* inputs, void* const* outputs, 
    void* workspace, cudaStream_t stream) noexcept
{
    /*
     * Plugin的核心的地方。每个插件都有一个自己的定制方案
     * Plugin直接调用kernel的地方
    */
    int nElements = 1;
    for (int i = 0; i < inputDesc[0].dims.nbDims; i++){
        nElements *= inputDesc[0].dims.d[i];
    }

    customScalarImpl(
            static_cast<const float*>(inputs[0]),
            static_cast<float*>(outputs[0]), 
            mParams.scalar, 
            mParams.scale,
            nElements,
            stream);

    return 0;
}

该函数的主要目的是处理传入的输入数据，通过调用一个 CUDA 核函数启动函数 customScalarImpl 来应用插件定义的计算，并将结果输出，这通常在模型的推理阶段调用，涉及到具体的 GPU 计算：(from ChatGPT)

• 1. 输入和输出描述符
  • inputDesc 和 outputDesc 提供了输入和输出张量的描述信息，包括数据类型、维度信息等
• 2. 输入和输出数据：
  • inputs 和 outputs 是指向输入和输出数据的指针数组。这些数据在 GPU 上进行处理
• 3. 工作空间和 CUDA 流
  • workspace 提供了用于计算的临时存储空间
  • stream 是 CUDA 流，用于异步执行 GPU 操作，确保操作的正确并发执行
• 4. 计算总元素数量
  • 计算输入张量中的总元素数（nElements），这是通过遍历输入张量维度并计算所有维度大小的乘积完成的，这个值决定了 CUDA 核函数需要处理的数据量
• 5. CUDA 核函数调用
  • customScalarImpl 是一个 CUDA 核函数启动函数，它具体调用了实现插件定义的操作的核函数，该函数需要以下参数：
    • 输入数据（inputs[0]）和输出数据（outputs[0]），这些数据通过 static_cast 转换为正确的数据类型
    • mParams.scalar 和 mParams.scale，这是插件特定的参数，控制具体的计算逻辑
    • nElements，告诉核函数需要处理的元素总数
    • stream，CUDA 流，用于管理 CUDA 操作的执行

customScalarImpl 的实现如下：

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

__global__ void customScalarKernel(
    const float* input, float* output, 
    const float scalar, const float scale, const int nElements)
{
    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= nElements) 
        return;

    output[index] = (input[index] + scalar) * scale;
}

void customScalarImpl(const float* inputs, float* outputs, const float scalar, const float scale, const int nElements, cudaStream_t stream)
{
    dim3 blockSize(256, 1, 1);
    dim3 gridSize(ceil(float(nElements) / 256), 1, 1);
    customScalarKernel<<<gridSize, blockSize, 0, stream>>>(inputs, outputs, scalar, scale, nElements);
}

customScalarImpl 函数与其对应的 CUDA 核函数 customScalarKernel 一起实现了在 GPU 上对输入张量进行简单的数学变换，这个过程涉及向每个元素添加一个标量值，然后乘以一个缩放因子

我们在来看下其中的 supportsFormatCombination 函数，其实现如下：

bool CustomScalarPlugin::supportsFormatCombination(int32_t pos, const PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
    /* 
     * 设置这个Plugin支持的Datatype以及TensorFormat, 每个插件都有自己的定制
     * 作为案例展示，这个customScalar插件只支持FP32，如果需要扩展到FP16以及INT8，需要在这里设置
    */
    
    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;
}

该函数用于指定插件支持的数据类型和张量格式组合，这是 TensorRT 插件开发中重要的一部分，确保插件能够处理指定格式的数据，值得注意的是 TensorFormat 数据格式有多种：

Layout Format for CHW

Layout Format for HWC

更多细节大家可以参考：https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#appendix

3. LeakyReLU Plugin案例

在 main 函数中还提供了另一个插件案例 Leaky ReLU，那它其实和上面的 scalar 差不多，我们简单过一下和前面的有什么差异：

CustomLeakyReLUPlugin::CustomLeakyReLUPlugin(const std::string &name, float alpha):
    mName(name)
{
    mParams.alpha = alpha;
    if (alpha < 0.0F) LOGE("ERROR detected when initialize plugin");
}

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<const float*>(inputs[0]),
            static_cast<float*>(outputs[0]), 
            mParams.alpha, 
            nElements,
            stream);

    return 0;
}

首先 Plugin 的构造函数和 enqueue 调用的核函数实现有所不同，其具体的核函数实现如下：

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

__global__ void customLeakyReLUKernel(
    const float* input, float* output, 
    const float alpha, const int nElements)
{
    const int index = blockIdx.x * blockDim.x + threadIdx.x;
    if (index >= nElements) 
        return;

    output[index] = input[index] > 0 ? input[index] : input[index] * alpha;
}

void customLeakyReLUImpl(const float* inputs, float* outputs, const float alpha, const int nElements, cudaStream_t stream)
{
    dim3 blockSize(256, 1, 1);
    dim3 gridSize(ceil(float(nElements) / 256), 1, 1);
    customLeakyReLUKernel<<<gridSize, blockSize, 0, stream>>>(inputs, outputs, alpha, nElements);
}

那此外很多的函数和前面的 scalar 都差不多，这边就不再分析了，所以大家能够发现插件中很多部分都是相同的，因此我们完全可以将它封装起来，这个我们之前在杜老师的课程中有提到过，大家感兴趣的可以看看：5.5.tensorRT基础(2)-封装插件过程，并实现更容易的插件开发

我们来执行下该案例，先看下 python 执行结果：

再对比看看 C++ 执行结果：

可以看到二者结果差不多，说明我们成功创建了 LeakyRelu 插件并完成了推理

总结

本次课程我们学习了 TensorRT 自定义插件的编写，主要是通过 scalar 和 leaky relu 两个案例来学习插件的整体流程，首先我们需要在 Python 中通过 symbolic 方法导出自定义的算子，接着我们需要在 C++ 上解析该算子并完成推理。我们需要定义两个类，一个是 Plugin 继承自 IPluginV2DynamicExt，另一个是 PluginCreator 继承自 IPluginCreator，然后分别实现其中的构造函数和虚函数，真正推理的函数 enqueue 是通过调用其相应的核函数来完成的，因此大家如果对流程熟悉了，本质上其实就是写自定义算子的 CUDA 核函数实现。

OK，以上就是 5.7 小节案例的全部内容了，下节我们来学习 5.8 小节插件的单元测试，敬请期待😄

下载链接

• tensorrt_starter源码
• 5.7-custom-basic-trt-plugin案例文件

参考

• Ubuntu20.04软件安装大全
• https://github.com/kalfazed/tensorrt_starter.git
• 5.5.tensorRT基础(2)-封装插件过程，并实现更容易的插件开发
• https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#appendix