乐趣区

关于深度学习:实现TensorRT70插件自由如果不踩坑使用TensorRT插件功能

本系列为新 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 版本的IPluginV2IOExtIPluginV2DynamicExt。将来不晓得会不会进去新的 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.cppcustom.h,copy 并 paste 官网代码,名字替换成本人的。以最新的IPluginV2DynamicExt 类为接口。

咱们须要写两个类:

  • MyCustomPlugin,继承IPluginV2DynamicExt,是插件类,用于写插件具体的实现
  • MyCustomPluginCreator,继承BaseCreator,是插件工厂类,用于依据需要创立该插件

对了,插件类继承 IPluginV2DynamicExt 才能够反对动静尺寸,其余插件类接口例如 IPluginV2IOExt 和前者大部分是类似的。

// 继承 IPluginV2DynamicExt 就够啦
class MyCustomPlugin final : public nvinfer1::IPluginV2DynamicExt

class MyCustomPluginCreator : public BaseCreator

MyCustomPlugin 插件类

总览:

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);

析构函数则须要执行 terminateterminate 函数就是开释这个 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);
}

上述步骤中,会提供 pluginNamepluginVersion初始化 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].formatinOut[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()->registerCreatorpluginCreator进行了注册,这样就实现注册工作了:

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″ 获取老潘学习路线材料与文章汇总,还有更多等你开掘。如果不想错过老潘的最新推文,请点击神秘链接。

退出移动版