本系列为新TensorRT的第一篇，为什么叫新，因为之前曾经写了两篇对于TensorRT的文章，是对于TensorRT-5.0版本的。好久没写对于TensorRT的文章了，所幸就以新来结尾吧~
接下来将要解说的TensorRT，将会是基于7.0版本。
7版本结尾的TensorRT变动还是挺大的，减少了很多新个性，然而TensorRT的外围运作形式还是没有什么变动的，对于TensorRT的介绍能够看之前写的这两篇：
- [利用TensorRT对深度学习进行减速
](https://oldpan.me/archives/us...
- [利用TensorRT实现神经网络提速(读取ONNX模型并运行)
](https://oldpan.me/archives/te...
本文的内容呢，次要是解说：
- TensorRT自定义插件的应用形式
- 如何增加本人的自定义算子
看完本篇能够让你少踩巨多坑，客官记得常来看啊。
前言
随着tensorRT的一直倒退(v5->v6->v7)，TensorRT的插件的应用形式也在不断更新。插件接口也在一直地变动，由v5版本的IPluginV2Ext，到v6版本的IPluginV2IOExt和IPluginV2DynamicExt。将来不晓得会不会进去新的API，不过这也不是咱要思考的问题，因为TensorRT的后兼容性做的很好，基本不必放心你写的旧版本插件在新版本上无奈运行。
目前的plugin-API：
TensorRT插件的存在目标，次要是为了让咱们实现TensorRT目前还不反对的算子，毕竟众口难调嘛，咱们在转换过程中必定会有op不反对的状况。这个时候就须要应用TensorRT的plugin去实现咱们的本人的op。此时咱们须要通过TensorRT提供的接口去实现本人的op，因而这个plugin的生命周期也须要遵循TensorRT的规定。
一个简略的理解
那么plugin到底长啥样，能够先看看TensorRT的官网plugin库长啥样，截止写这篇文章时，master分支是7.2版本的plugin：
https://github.com/NVIDIA/Ten...
官网提供的插件曾经相当多，而且TensorRT开源了plugin局部(能够让咱们白嫖！)。并且能够看到其源码，通过模拟源码来学习plugin是如何写的。
如果要增加本人的算子，能够在官网的plugin库里头进行批改增加，而后编译官网的plugin库。将生成的libnvinfer_plugin.so.7替换本来的.so文件即可。或者本人写一个相似于官网plugin的组件，将名称替换一下，同样生成.so，在TensorRT的推理我的项目中援用这个动态链接库即可。
以下介绍中，咱们须要写的IPlugin简称为插件op。
开始写插件
有趣味的能够先看看TensorRT的官网文档，官网文档的介绍简略意骇，不过坑是少不了的..而本文的目标，就是尽量让你少趟坑。
首先依照官网plugin的排布形式，上面轻易挑了个官网plugin：
筹备一个本人的插件:custom.cpp和custom.h，copy并paste官网代码，名字替换成本人的。以最新的IPluginV2DynamicExt类为接口。
咱们须要写两个类：
MyCustomPlugin，继承IPluginV2DynamicExt，是插件类，用于写插件具体的实现MyCustomPluginCreator，继承BaseCreator，是插件工厂类，用于依据需要创立该插件
对了，插件类继承IPluginV2DynamicExt才能够反对动静尺寸，其余插件类接口例如IPluginV2IOExt和前者大部分是类似的。
// 继承IPluginV2DynamicExt就够啦class MyCustomPlugin final : public nvinfer1::IPluginV2DynamicExtclass MyCustomPluginCreator : public BaseCreatorMyCustomPlugin 插件类
总览：
class MyCustomPlugin final : public nvinfer1::IPluginV2DynamicExt{public: MyCustomPlugin( int in_channel, const std::vector<float>& weight, const std::vector<float>& bias); MyCustomPlugin( int in_channel, nvinfer1::Weights const& weight, nvinfer1::Weights const& bias); MyCustomPlugin(void const* serialData, size_t serialLength); MyCustomPlugin() = delete; ~MyCustomPlugin() override; int getNbOutputs() const override; DimsExprs getOutputDimensions(int outputIndex, const nvinfer1::DimsExprs* inputs, int nbInputs, nvinfer1::IExprBuilder& exprBuilder) override; int initialize() override; void terminate() override; size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs, int nbInputs, const nvinfer1::PluginTensorDesc* outputs, int nbOutputs) const override; int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) override; size_t getSerializationSize() const override; void serialize(void* buffer) const override; bool supportsFormatCombination(int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs) override; const char* getPluginType() const override; const char* getPluginVersion() const override; void destroy() override; nvinfer1::IPluginV2DynamicExt* clone() const override; void setPluginNamespace(const char* pluginNamespace) override; const char* getPluginNamespace() const override; DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const override; void attachToContext(cudnnContext* cudnn, cublasContext* cublas, nvinfer1::IGpuAllocator* allocator) override; void detachFromContext() override; void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in, int nbInputs, const nvinfer1::DynamicPluginTensorDesc* out, int nbOutputs) override;private: int _in_channel; std::vector<float> weight; std::vector<float> bias; float* weight; float* bias; bool _initialized; const char* mPluginNamespace; std::string mNamespace;};成员变量
如果你的插件有weights(相似于conv操作的weight和bias)，有参数(相似于conv中的kernel-size、padding)，在类中则须要定义为成员变量，为private类型：
以MyCustomPlugin为例，假如咱们的这个MyCustomPlugin有两个权重weight和bias以及一个参数in_channel(这个权重和参数没有啥意义，纯正，纯正为了演示)：
private: int _in_channel; // 参数 std::vector<float> _weight; // 权重，在cpu空间寄存 std::vector<float> _bias; // 偏置权重，在cpu空间寄存 float* _d_weight; // 权重，在GPU空间寄存 float* _d_bias; bool _initialized; cudnnHandle_t _cudnn_handle; const char* mPluginNamespace; std::string mNamespace;构造函数和析构函数
构造函数个别设置为三个。
第一个用于在parse阶段，PluginCreator用于创立该插件时调用的构造函数，须要传递权重信息以及参数。
第二个用于在clone阶段，复制这个plugin时会用到的构造函数。
第三个用于在deserialize阶段，用于将序列化好的权重和参数传入该plugin并创立爱你哦。
以咱们的MyCustomPlugin为例：
MyCustomPlugin(int in_channel, nvinfer1::Weights const& weight, nvinfer1::Weights const& bias);MyCustomPlugin(float in_channel, const std::vector<float>& weight, const std::vector<float>& bias);MyCustomPlugin(void const* serialData, size_t serialLength);析构函数则须要执行terminate，terminate函数就是开释这个op之前开拓的一些显存空间:
MyCustomPlugin::~MyCustomPlugin(){ terminate();}留神须要把默认构造函数删掉：
MyCustomPlugin() = delete;getNbOutputs
插件op返回多少个Tensor，比方MyCustomPlugin这个操作只输入一个Tensor(也就是一个output)，所以间接return 1：
// MyCustomPlugin returns one output.int MyCustomPlugin::getNbOutputs() const{ return 1;}initialize
初始化函数，在这个插件筹备开始run之前执行。
次要初始化一些提前开拓空间的参数，个别是一些cuda操作须要的参数(例如conv操作须要执行卷积操作，咱们就须要提前开拓weight和bias的显存)，如果咱们的算子须要这些参数，则在这里须要提前开拓显存。
须要留神的是，如果插件算子须要开拓比拟大的显存空间，不倡议本人去申请显存空间，能够应用Tensorrt官网接口传过来的workspace指针来获取显存空间。因为如果这个插件被一个网络调用了很屡次，而这个插件op须要开拓很多显存空间，那么TensorRT在构建network的时候会依据这个插件被调用的次数开拓很多显存，很容易导致显存溢出。
getOutputDataType
返回后果的类型，一般来说咱们插件op返回后果类型与输出类型统一：
nvinfer1::DataType InstanceNormalizationPlugin::getOutputDataType( int index, const nvinfer1::DataType* inputTypes, int nbInputs) const{ ASSERT(inputTypes && nbInputs > 0 && index == 0); return inputTypes[0];}getWorkspaceSize
这个函数须要返回这个插件op须要两头显存变量的理论数据大小(bytesize)，这个是通过TensorRT的接口去获取，是比拟标准的形式。
咱们须要在这里确定这个op须要多大的显存空间去运行，在理论运行的时候就能够间接应用TensorRT开拓好的空间而不是本人去申请显存空间。
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);}enqueue
理论插件op的执行函数，咱们本人实现的cuda操作就放到这里(当然C++写的op也能够放进来，不过因为是CPU执行，速度就比较慢了)，与平常一样承受输出inputs产生输入outputs，传给相应的指针就能够。
int enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream){ // 如果这个fun是你须要的两头变量 这里能够间接用TensorRT为你开拓的显存空间 fun = static_cast<float*>(workspace); }须要留神的是，如果咱们的操作须要一些散布在显存中的两头变量，能够通过传过来的指针参数workspace获取，上述代码简略阐明了一下应用办法。
再多说一句，咱们默认写的.cu是fp32的，TensorRT在fp16运行模式下，运行到不反对fp16的插件op时，会主动切换到fp32模式，等插件op运行完再切换回来。
getOutputDimensions
TensorRT反对Dynamic-shape的时候，batch这一维度必须是explicit的，也就是说，TensorRT解决的维度从以往的三维[3,-1,-1]变成了[1,3,-1,-1]。最新的onnx-tensorrt也必须设置explicit的batchsize，而且这个batch维度在getOutputDimensions中是能够获取到的。
在旧版的IPluginV2类中，getOutputDimensions的定义如下：
virtual Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) TRTNOEXCEPT = 0;而在新版的IPluginV2DynamicExt类中定义如下：
virtual DimsExprs getOutputDimensions(int outputIndex, const DimsExprs* inputs, int nbInputs, IExprBuilder& exprBuilder) = 0;咱们要做的就是在这个成员函数中依据输出维度推理出模型的输入维度，须要留神的是，尽管说输入维度
是由输出维度决定，但这个输入维度其实“内定”的(也就是在计算之前就算进去了)。如果咱的插件op的输入维度须要通过理论运行计算失去，那么这个函数就无奈满足咱了。
set/getPluginNamespace
为这个插件设置namespace名字，如果不设置则默认是""，须要留神的是同一个namespace下的plugin如果名字雷同会抵触。
PluginFieldCollection
这个是成员变量，也会作为getFieldNames成员函数的返回类型。PluginFieldCollection的次要作用是传递这个插件op所须要的权重和参数，在理论的engine推理过程中并不应用，而在parse中会用到(例如caffe2trt、onnx2trt)。
当应用这些parse去解析这个op的时候，这个op的权重和参数会经验Models --> TensorRT engine --> TensorRT runtime这个过程。
举个例子，在onnx-tensorrt中，咱们用过DEFINE_BUILTIN_OP_IMPORTER去注册op，而后通过parse解析onnx模型，依据注册好的op去一个个解析构建模型，如果咱们定义的op为my_custom_op,在DEFINE_BUILTIN_OP_IMPORTER(my_custom_op)会这样实现：
DEFINE_BUILTIN_OP_IMPORTER(mycustom_op){ ASSERT(inputs.at(0).is_tensor(), ErrorCode::kUNSUPPORTED_NODE); ... const std::string pluginName = "CUSTOM-OP"; const std::string pluginVersion = "001"; // 这个f保留这个op须要的权重和参数，从onnx模型中获取 std::vector<nvinfer1::PluginField> f; f.emplace_back("in_channel", &in_channel, nvinfer1::PluginFieldType::kINT32, 1); f.emplace_back("weight", kernel_weights.values, nvinfer1::PluginFieldType::kFLOAT32, kernel_weights.count()); f.emplace_back("bias", bias_weights.values, nvinfer1::PluginFieldType::kFLOAT32, bias_weights.count); // 这个从将plugin工厂中获取该插件，并且将权重和参数传递进去 nvinfer1::IPluginV2* plugin = importPluginFromRegistry(ctx, pluginName, pluginVersion, node.name(), f); RETURN_FIRST_OUTPUT(ctx->network()->addPluginV2(tensors.data(), tensors.size(), *plugin));}进入importPluginFromRegistry函数外部，能够发现参数通过fc变量通过createPlugin传递给了plugin：
nvinfer1::IPluginV2* importPluginFromRegistry(IImporterContext* ctx, const std::string& pluginName, const std::string& pluginVersion, const std::string& nodeName, const std::vector<nvinfer1::PluginField>& pluginFields){ const auto mPluginRegistry = getPluginRegistry(); const auto pluginCreator = mPluginRegistry->getPluginCreator(pluginName.c_str(), pluginVersion.c_str(), "ONNXTRT_NAMESPACE"); if (!pluginCreator) { return nullptr; } // 承受传进来的权重和参数信息 传递给plugin nvinfer1::PluginFieldCollection fc; fc.nbFields = pluginFields.size(); fc.fields = pluginFields.data(); return pluginCreator->createPlugin(nodeName.c_str(), &fc);}上述步骤中，会提供pluginName和pluginVersion初始化MyCustomPluginCreator，其中createPlugin成员函数是咱们须要编写的(下文会说)。
configurePlugin
配置这个插件op，判断输出和输入类型数量是否正确。官网还提到通过这个配置信息能够告知TensorRT去抉择适合的算法(algorithm)去调优这个模型。
但主动调优目前还没有尝试过，咱们个别本人写的plugin执行代码都是定死的，所谓的调优步骤可能更多地针对官网的op。
上面的plugin中configurePlugin函数仅仅是简略地确认了下输出和输入以及类型。
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);}clone
这玩意儿干嘛的，顾名思义，就是克隆嘛，将这个plugin对象克隆一份给TensorRT的builder、network或者engine。这个成员函数会调用上述说到的第二个构造函数：
MyCustomPlugin(float in_channel, const std::vector<float>& weight, const std::vector<float>& bias);将要克隆的plugin的权重和参数传递给这个构造函数。
IPluginV2DynamicExt* MyCustomPlugin::clone() const{ // auto plugin = new MyCustomPlugin{_in_channel, _weight, _bias}; plugin->setPluginNamespace(mPluginNamespace); return plugin;}clone成员函数次要用于传递不变的权重和参数，将plugin复制n多份，从而能够被不同engine或者builder或者network应用。
getSerializationSize
返回序列化时须要写多少字节到buffer中。
size_t MyCustomPlugin::getSerializationSize() const{ return (serialized_size(_in_channel) + serialized_size(_weight) + serialized_size(_bias) );}supportsFormatCombination
TensorRT调用此办法以判断pos索引的输出/输入是否反对inOut[pos].format和inOut[pos].type指定的格局/数据类型。
如果插件反对inOut[pos]处的格局/数据类型，则返回true。 如果是否反对取决于其余的输出/输入格局/数据类型，则插件能够使其后果取决于inOut[0..pos-1]中的格局/数据类型，该格局/数据类型将设置为插件反对的值。 这个函数不须要查看inOut[pos + 1..nbInputs + nbOutputs-1]，pos的决定必须仅基于inOut[0..pos]。
bool MyCustomPlugin::supportsFormatCombination( int pos, const nvinfer1::PluginTensorDesc* inOut, int nbInputs, int nbOutputs){ // 假如有一个输出一个输入 assert(0 <= pos && pos < 2); const auto *in = inOut; const auto *out = inOut + nbInputs; switch (pos) { case 0: return in[0].type == DataType::kFLOAT && in[0].format == nvinfer1::TensorFormat::kLINEAR; case 1: return out[0].type == in[0].type && out[0].format == nvinfer1::TensorFormat::kLINEAR; }}serialize
把须要用的数据依照程序序列化到buffer外头。
void MyCustomPlugin::serialize(void *buffer) const{ serialize_value(&buffer, _in_channel); serialize_value(&buffer, _weight); serialize_value(&buffer, _bias);}attachToContext
如果这个op应用到了一些其余货色，例如cublas handle，能够间接借助TensorRT外部提供的cublas handle:
void MyCustomPlugin::attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator){ mCublas = cublasContext;}MyCustomPluginCreator 插件工厂类
总览：
class MyCustomPluginCreator : public BaseCreator{public: MyCustomPluginCreator(); ~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; static std::vector<PluginField> mPluginAttributes; std::string mNamespace;};构造函数
创立一个空的mPluginAttributes初始化mFC。
MyCustomPluginCreator::MyCustomPluginCreator(){ mPluginAttributes.emplace_back(PluginField("in_channel", nullptr, PluginFieldType::kFLOAT32, 1)); mPluginAttributes.emplace_back(PluginField("weight", nullptr, PluginFieldType::kFLOAT32, 1)); mPluginAttributes.emplace_back(PluginField("bias", nullptr, PluginFieldType::kFLOAT32, 1)); mFC.nbFields = mPluginAttributes.size(); mFC.fields = mPluginAttributes.data();}createPlugin
这个成员函数作用是通过PluginFieldCollection去创立plugin，将op须要的权重和参数一个一个取出来，而后调用上文提到的第一个构造函数：
MyCustomPlugin(int in_channel, nvinfer1::Weights const& weight, nvinfer1::Weights const& bias);去创立plugin。
MyCustomPlugin示例：
IPluginV2DynamicExt* MyCustomPlugin::createPlugin(const char* name, const nvinfer1::PluginFieldCollection* fc){ int in_channel; std::vector<float> weight; std::vector<float> 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<const int32_t*>(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<const float*>(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<const float*>(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;}deserializePlugin
这个函数会被onnx-tensorrt的一个叫做TRT_PluginV2的转换op调用，这个op会读取onnx模型的data数据将其反序列化到network中。
一些官网插件的注意事项
应用官网插件会遇到些小问题。
topk问题
官网的topk插件最多反对k<=3840。否则会报：
[TensorRT] ERROR: Parameter check failed at: ../builder/Layers.cpp::TopKLayer::3137, condition: k > 0 && k <= MAX_TOPK_K
相干问题：https://github.com/tensorflow...
batchednms问题
官网的batchednms最大反对的topk为4096，太大也会解体。不过能够批改源代码实现冲破这个数值，但依然有bug：
void (*kernel[])(const int, const int, const int, const int, const float, const bool, const bool, float *, T_SCORE *, int *, T_SCORE *, int *, bool) = { P(1), P(2), P(3), P(4), P(5), P(6), P(7), P(8), P(9), P(10), P(11), P(12), P(13), P(14), P(15), P(16) };对于plugin的注册
简略说下plugin的注册流程。
在加载NvInferRuntimeCommon.h头文件的时候会失去一个getPluginRegistry，这里类中蕴含了所有曾经注册了的IPluginCreator，在应用的时候咱们通过getPluginCreator函数失去相应的IPluginCreator。
注册插件有两种形式，第一种能够看官网的plugin代码。
extern "C" {bool initLibNvInferPlugins(void* logger, const char* libNamespace){ initializePlugin<nvinfer1::plugin::GridAnchorPluginCreator>(logger, libNamespace); initializePlugin<nvinfer1::plugin::NMSPluginCreator>(logger, libNamespace); initializePlugin<nvinfer1::plugin::ReorgPluginCreator>(logger, libNamespace); ... return true;}其中initializePlugin函数执行了addPluginCreator函数：
template <typename CreatorType>void initializePlugin(void* logger, const char* libNamespace){ PluginCreatorRegistry::getInstance().addPluginCreator<CreatorType>(logger, libNamespace);}addPluginCreator函数又执行了getPluginRegistry()->registerCreator对pluginCreator进行了注册，这样就实现注册工作了：
void addPluginCreator(void* logger, const char* libNamespace){ ... if (mRegistryList.find(pluginType) == mRegistryList.end()) { bool status = getPluginRegistry()->registerCreator(*pluginCreator, libNamespace); if (status) { mRegistry.push(std::move(pluginCreator)); mRegistryList.insert(pluginType); verboseMsg = "Plugin creator registration succeeded - " + pluginType; } else { errorMsg = "Could not register plugin creator: " + pluginType; } } else { verboseMsg = "Plugin creator already registered - " + pluginType; } ...}另一种注册能够间接通过REGISTER_TENSORRT_PLUGIN来注册：
//!//! \brief Return the plugin registry//!// 在加载`NvInferRuntimeCommon.h`头文件的时候会失去一个`getPluginRegistry`extern "C" TENSORRTAPI nvinfer1::IPluginRegistry* getPluginRegistry();namespace nvinfer1{template <typename T>class PluginRegistrar{public: PluginRegistrar() { getPluginRegistry()->registerCreator(instance, ""); }private: T instance{};};#define REGISTER_TENSORRT_PLUGIN(name) \ static nvinfer1::PluginRegistrar<name> pluginRegistrar##name {}} // namespace nvinfer1也就是说，如果咱们曾经在plugin的.h文件中执行了REGISTER_TENSORRT_PLUGIN(BatchedNMSPluginCreator);就不须要再创立一个相似于官网的initLibNvInferPlugins()函数去一个一个注册了。
参考链接
https://github.com/NVIDIA/Ten...
https://github.com/triton-inf...
https://blog.csdn.net/u010552...
https://docs.nvidia.com/deepl...
https://forums.developer.nvid...
https://forums.developer.nvid...
https://github.com/NVIDIA/Ten...
https://forums.developer.nvid...
DCNv2-github
https://github.com/CharlesSha...
https://github.com/chengdazhi...
交换
如果你与我气味相投于此，老潘很违心与你交换；如果你喜爱老潘的内容，欢送关注和反对。博客每周更新一篇深度原创文，关注公众号「oldpan博客」不错过最新文章。老潘也会整顿一些本人的私藏，心愿能帮忙到大家，公众号回复"888"获取老潘学习路线材料与文章汇总，还有更多等你开掘。如果不想错过老潘的最新推文，请点击神秘链接。