TensorRT7自定义插件实现教程

前言

随着tensorRT的不断发展(v5->v6->v7)，TensorRT的插件的使用方式也在不断更新。插件接口也在不断地变化，由v5版本的IPluginV2Ext，到v6版本的IPluginV2IOExt和IPluginV2DynamicExt。未来不知道会不会出来新的API，不过这也不是咱要考虑的问题，因为TensorRT的后兼容性做的很好，一般两三个大版本间不用担心你写的旧版本插件在新版本上无法运行，不过还是要多关注下版本发布说明里的plugin改变，如TensorRT723 Plugin API说明。

目前的plugin-API：

TensorRT目前只支持一些常见的操作，有很多操作它并不支持，这时候就需要我们使用TensorRT插件接口自行编写TensorRT的插件层，从而使得这些不能支持的操作能在tensorRT中使用，同时这个plugin的生命周期也需要遵循TensorRT的规则。

NVIDIA也提供了部分插件并将其开源了，如TensorRT7.2插件表，我们可以模仿源码来编写自己的插件。引入自己的算子有两种方法：

• 在官方的plugin库里进行修改添加，然后编译官方plugin库生成动态链接文件libnvinfer_plugin.so.7替换原本的动态链接文件。可以参考文末"补充1"。

• 按照其中样例编写一个自己的组件，编译生成单独的动态链接文件，在项目中引用该库。本文以此方式为例。

编写插件

本文以伪代码形式为例，说明插件编写。使用创建以下两个文件，文件名随意，使用cuda C语言。

yourlayer.h:

namespace nvinfer1
{
  class API YourPlugin: public IPluginV2IOExt
  {
  public:
      //一些必要方法
  private:
      //一些必要私有变量
  };
  class API YourPluginCreator: public IpluginCreator
  {
  public:
      //一些必要方法
  private:
      //一些必要私有变量
  };
  REGISTER_TENSORRT_PLUGIN(YourPluginCreator);
}

yourlayer.cu:

#include 
#include "yourlayer.h"

namespace nvinfer1
{
  //此构造函数用于传参数构建插件
  YourPlugin::YourPlugin() //此处笔者省略，可自行添加需要传入析构函数的参数
  {
      //实现
  }
  //此构造函数用于runtime从序列化后的比特流构建插件
  YourPlugin::YourPlugin(const void* data, size_t length)
  {
      //实现
  }
  
  YoloLayerPlugin::~YoloLayerPlugin()  //析构函数
  {
      //实现
  }
  
  ...  //只是展示大概结构，所以笔者省略部分IPluginV2IOExt类方法。
  
  PluginFieldCollection YourPluginCreator::mFC{};
  std::vector<PluginField> YourPluginCreator::mPluginAttributes;
  
  IPluginV2IOExt* YourPluginCreator::createPlugin()
  {
  }
  
  ... //只是展示大概结构，所以笔者省略部分Creator类方法。
}

插件类

本部分描述了IPluginV2IOExt类。为了将插件层连接到相邻层并设置输入和输出数据结构，builder通过调用以下插件方法检查输出数量及其维度。以下代码均声明于IPluginV2IOExt类的大括号内。

2(或3)个构造函数和1个析构函数

YourPlugin();//此处笔者省略参数
YourPlugin(const void* data, size_t length);
// It makes no sense to construct UffPoolPluginV2 without arguments.
YourPlugin() = delete;  //删除默认构造函数
~YourPlugin() {terminate();}

如果你是使用解析器导入模型，插件里还需要一个用于parser阶段的构造函数：

YourPlugin(int in_channel, nvinfer1::Weights const& weight, nvinfer1::Weights const& bias);//参数看你自己情况修改添加

getNbOutputs

指定输出张量数量。

int getNbOutputs() const override;

getOutputDimensions

根据输入维度推理出模型的输出维度。TensorRT7以后要求输入是全维度的，batch这一维度必须是explicit的，也就是说TensorRT处理的维度从以往的三维[3,-1,-1]变成了[1,3,-1,-1]。最新的onnx-tensorrt也必须设置explicit的batchsize，而且这个batch维度在getOutputDimensions中是可以获取到的。

Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) override;

注意： 虽然说输出维度是由输入维度决定，但这个输出维度其实“内定”的(也就是在计算之前就算出来了)。如果你的插件的输出维度需要通过实际运行计算得到，那么这个函数就无法满足。

supportsFormatCombination

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

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

getOutputDataType

用于获取给定索引处输出的数据类型。返回的数据类型必须具有插件支持的格式.

插件层可以支持四种数据格式和布局，例如：

• NCHW 单 (FP32)、半精度 (FP16) 和整数 (INT32) 张量
• NC/2HW2 和 NHWC8 半精度 (FP16) 张量

格式被PluginFormatType枚举.

不计算所有数据并且除了输入和输出张量还需要内存空间的插件可以通过getWorkspaceSize指定额外的获取的内存要求，由builder调用以确定和预分配暂存空间。

DataType getOutputDataType(int index, const DataType* inputTypes, int nbInputs) const override;

在构建和推理期间，插件层可能被多次配置和执行。在构建时，为了发现最佳配置，层被配置、初始化、执行和终止。一旦为插件选择了最佳格式，插件将被再次配置，然后它会被初始化一次并在推理应用程序的生命周期内根据需要执行多次，最终在引擎销毁时终止。这些步骤由builder和engine使用以下插件方法控制：

configurePlugin

传递输入和输出的数量、所有输入和输出的维度和数据类型、所有输入和输出的广播信息、所选插件格式和maximum batch size。此时，插件设置其内部状态，并为给定配置选择最合适的算法和数据结构。

virtual void configurePlugin(const PluginTensorDesc* in, int32_t nbInput, const PluginTensorDesc* out, int32_t nbOutput) override;

initialize

为execution初始化层。初始化前配置已知并且推理enine正在被创建，所以插件可以设置它的内部数据结构。主要初始化一些提前开辟空间的参数，一般是一些cuda操作需要的参数(例如conv操作需要执行卷积操作，我们就需要提前开辟weight和bias的显存)，假如我们的算子需要这些参数，则在这里需要提前开辟显存。

int initialize() override;

注意： 如果插件算子需要开辟比较大的显存空间，不建议自己去申请显存空间，可以使用Tensorrt官方接口传过来的workspace指针来获取显存空间。因为如果这个插件被一个网络调用了很多次，而这个插件op需要开辟很多显存空间，那么TensorRT在构建network的时候会根据这个插件被调用的次数开辟很多显存，很容易导致显存溢出。

enqueue

封装插件的实际算法和内核调用，并提供runtime batch size、指向输入、输出和暂存空间的指针，以及用于内核执行的CUDA流。

virtual int enqueue(int32_t batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override;
{
 // 假如这个fun是你需要的中间变量 这里可以直接用TensorRT为你开辟的显存空间
 fun  = static_cast<float*>(workspace);
}

注意： 如果我们的操作需要一些分布在显存中的中间变量，可以通过传过来的指针参数workspace获取，下面代码简单说明了一下使用方法。另外，我们默认写的.cu是fp32的，TensorRT在fp16运行模式下，运行到不支持fp16的插件op时，会自动切换到fp32模式，等插件op运行完再切换回来。

int YourPlugin::enqueue(int32_t batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) override
{
 // 假如这个fun是你需要的中间变量 这里可以直接用TensorRT为你开辟的显存空间
 fun  = static_cast<float*>(workspace);
 forwardGpu()
}

terminate

engine context被销毁时调用，释放插件层初始化期间被获取的资源。

virtual void terminate() override;

clone

每次创建包含此插件层的新构建器、网络或引擎时都会调用此方法。它应该返回一个具有正确参数的新插件对象

IPluginV2IOExt* clone() const override;

destroy

用于销毁插件对象。每当builder或network或engine被销毁时都会调用它

void destroy() override {delete this;};

set/getPluginNamespace

该方法用于设置该插件对象所属的库命名空间（默认可以是""）。来自同一个插件库的所有插件对象应该具有相同的命名空间.

void setPluginNamespace(const char* libNamespace) override;
const char* getPluginNamespace() const override;

getPluginType/Version

返回插件类型和版本。应与相应插件创建者返回的插件名称/版本匹配。

const char* getPluginType() const override;
const char* getPluginVersion() const override;

getWorkspaceSize/getSerializationSize

这个函数需要返回这个插件op需要中间显存变量的实际数据大小(bytesize)，这个是通过TensorRT的接口去获取，是比较规范的方式。我们需要在这里确定这个op需要多大的显存空间去运行，在实际运行的时候就可以直接使用TensorRT开辟好的空间而不是自己去申请显存空间。

获取序列化需要的buffer size。

virtual size_t getWorkspaceSize(int maxBatchSize) const override;
virtual size_t getSerializationSize() const override;

serialize

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

virtual void serialize(void* buffer) const override;

attachToContext/detachFromContext

如果这个op使用到了一些其他东西，例如cublas handle，可以直接将插件对象附加到执行上下文并授予插件访问某些上下文资源的权限,借助TensorRT内部提供的cublas handle.
从它的执行上下文中分离插件对象.
疑问: 看了若干例子，这两个方法貌似不是必要的？

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

IPluginV2IOExt支持可以处理广播输入和输出的插件。此功能需要实现以下方法：

canBroadcastInputAcrossBatch

为每个输入调用此方法，其张量在一个批次中语义广播。如果 canBroadcastInputAcrossBatch返回true（意味着插件可以支持广播），TensorRT不会复制输入张量。插件应该在批处理中共享一个副本。如果它返回false, TensorRT将复制输入张量，使其看起来像一个非广播张量。

bool canBroadcastInputAcrossBatch(int inputIndex) const override;

isOutputBroadcastAcrossBatch

这是为每个输出索引调用的。插件应该在给定索引处的输出返回true并在批处理中广播。

bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const override;

插件创建器类

IPluginCreator类中以下方法用于从插件注册表中查找和创建适当的插件。

1个构造函数和1个析构函数

构造与析构函数声明

YourPluginCreator();

~YourPluginCreator() override = default;

构造函数实现：初始化mFC

//创建一个空的mPluginAttributes在构造函数里初始化mFC
PluginFieldCollection YourPluginCreator::mFC{};
std::vector<PluginField> YourPluginCreator::mPluginAttributes;

YourPluginCreator::YourPluginCreator()
 {
     mPluginAttributes.clear();

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

getPluginName/Version

这将返回插件名称，应与IPluginExt::getPluginType返回值匹配。
返回插件版本，对于所有内部TensorRT插件，默认为1。

const char* getPluginName() const override;
const char* getPluginVersion() const override;

getFieldNames

为了成功创建插件，需要知道插件的所有字段参数。此方法返回带有PluginField条目的PluginFieldCollection结构，这个条目被填充以反映字段的name和PluginFieldType（数据应该指向nullptr）

const PluginFieldCollection* getFieldNames() override;

createPlugin

此方法用于使用PluginFieldCollection参数创建插件。PluginField条目的data字段应被填充为指向每个插件字段条目的实际数据。

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

deserializePlugin

此方法由TensorRT引擎根据插件名称和版本在内部调用。它返回用于推理的插件对象。

IPluginV2IOExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;

set/getPluginNamespace

该方法用于设置这个创建器实例所属的命名空间（默认可以是“”）。

void setPluginNamespace(const char* libNamespace) override;
const char* getPluginNamespace() const override;

编译插件

CMakeLists.txt:

cmake_minimum_required(VERSION 2.6)

add_definitions(-std=c++11)
add_definitions(-DAPI_EXPORTS)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)

if(WIN32)
enable_language(CUDA)
endif(WIN32)

# cuda和tensorrt的include和link目录, 如果你的不同需要自己修改
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/local/tensorrt/include/)
link_directories(/usr/local/tensorrt/lib/)

#-g是dubug 模式------------------这个非常重要，否则无法进入断点调试
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")

# yourplugin为你编译出的动态库名，自己修改
cuda_add_library(yourplugin SHARED ${PROJECT_SOURCE_DIR}/yourlayer.cu)
target_link_libraries(yourplugin nvinfer cudart)

if(UNIX)
add_definitions(-O2 -pthread)
endif(UNIX)

编译生成libyourplugin.so:

$ mkdir build && cd build
$ cmake ..
$ make -j4

添加插件层

添加插件

使用插件创建器，IPluginCreator::createPlugin()可以被调用，返回类型为IPluginV2的插件对象. 可以使用addPluginV2()将此对象添加到TensorRT网络，它创建并添加一个层到网络，然后将该层绑定到给定的插件。该方法还返回一个指向层的指针（类型为IPluginV2Layer)，可用于访问层或插件本身（通过getPlugin()）。

例如，将插件层添加到你的网络，插件名称设置为pluginName和版本设置为pluginVersion：

C++:

//使用外部函数getPluginRegistry获取全局TensorRT插件注册库
auto creator = getPluginRegistry()->getPluginCreator(pluginType, pluginVersion);

//如果使用TRT内部已注册的插件就直接使用下方代码
//const PluginFieldCollection* pluginFC = creator->getFieldNames();
//填充插件层的字段参数layerFields 
//PluginFieldCollection *pluginData = parseAndFillFields(pluginFC, layerFields);

//如果使用自己编写的插件，需要自行编写插件需要的PluginField填充到PluginFieldCollection
...//笔者省略编写PluginField
PluginFieldCollection plugin_data；

//使用layerName和插件元数据创建插件
IPluginV2 *pluginObj = creator->createPlugin(layerName, pluginData);

//使用network API向网络添加插件层
auto layer = network.addPluginV2(&inputs[0], int(inputs.size()), pluginObj);
… (build rest of the network and serialize engine)
pluginObj->destroy() // Destroy the plugin object
… (destroy network, engine, builder)
… (free allocated pluginData)

Python:

import tensorrt as trt
import numpy as np
import ctypes

ctypes.CDLL('libyourplugin.so')
TRT_LOGGER = trt.Logger()

#根据namespace初始化并注册插件，编写的插件和内置插件同属''的话，就全部初始化。
trt.init_libnvinfer_plugins(TRT_LOGGER, '')

# 如果创建多个插件可通过plugin_creator_list方式一次性创建
# PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list
# def get_trt_plugin(plugin_name):
#    plugin = None
#    for plugin_creator in PLUGIN_CREATORS:
#         if plugin_creator.name == plugin_name:
#            plugin_field = trt.PluginField("scale", np.array([0.1], dtype=np.float32), trt.PluginFieldType.FLOAT32)
#            field_collection = trt.PluginFieldCollection([plugin_field])
#            plugin = plugin_creator.create_plugin(name=plugin_name, field_collection=field_collection)
#    return plugin

#这里YourLayer_TRT与YourPlugin::getPluginType()返回值一致
plugin_creator = trt.get_plugin_registry().get_plugin_creator('YourLayer_TRT', "1")

#如果使用自己编写的插件，需要自行编写插件需要的PluginField填充到PluginFieldCollection
...#笔者省略编写PluginField
fc = trt.PluginFieldCollection()
plugin = plugin_creator.create_plugin(name="yourlayer",field_collection=fc)

def main():
    with trt.Builder(TRT_LOGGER) as builder, builder.create_network() as network:
        builder.max_workspace_size = 2**20
        input_layer = network.add_input(name="input_layer", dtype=trt.float32, shape=(1, 1))
        yourlayer = network.add_plugin_v2(inputs=[input_layer], plugin=plugin)
        yourlayer.get_output(0).name = "outputs"
        network.mark_output(yourlayer.get_output(0))

编译项目

CMakeLists.txt:

cmake_minimum_required(VERSION 2.6)
#项目名自己定
project(yourModel)

add_definitions(-std=c++11)
add_definitions(-DAPI_EXPORTS)
option(CUDA_USE_STATIC_CUDA_RUNTIME OFF)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_BUILD_TYPE Debug)

find_package(CUDA REQUIRED)

if(WIN32)
enable_language(CUDA)
endif(WIN32)

include_directories(${PROJECT_SOURCE_DIR}/include)
# cuda和tensorrt的include和link目录, 如果你的不同需要自己修改
# cuda
include_directories(/usr/local/cuda/include)
link_directories(/usr/local/cuda/lib64)
# tensorrt
include_directories(/usr/local/tensorrt/include/)
link_directories(/usr/local/tensorrt/lib/)

#-g是dubug 模式------------------这个非常重要，否则无法进入断点调试
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -g -std=c++11 -Wall -Ofast -Wfatal-errors -D_MWAITXINTRIN_H_INCLUDED")

# yourplugin为你编译出的动态库名，自己修改
cuda_add_library(yourplugin SHARED ${PROJECT_SOURCE_DIR}/yololayer.cu)
target_link_libraries(yourplugin nvinfer cudart)

#如果C++文件中用到opencv需要包含opencv库
find_package(OpenCV)
include_directories(${OpenCV_INCLUDE_DIRS})

target_link_libraries(yourModel nvinfer)
target_link_libraries(yourModel cudart)
target_link_libraries(yourModel yourplugin)
target_link_libraries(yourModel ${OpenCV_LIBS})

if(UNIX)
add_definitions(-O2 -pthread)
endif(UNIX)

参考

1. TensorRT-7.x自定义插件详细指南——老潘
2. 为tensorRT添加自定义层——知乎ltpyuanshuai

补充

1. TensorRT7 OSS实现插件
2. TensorRT自定义算子_torch2trt