【模型加速】自定义TensorRT NMS3D插件(2)

添加自定义TensorRT插件

        TensorRT的NMS Plugin在官网已有开源，但不符合我本次实验的需求，所以还得再自定义TensorRT Plugin。随着tensorRT的不断发展(v5->v6->v7)，TensorRT的插件的使用方式也在不断更新。插件接口也在不断地变化，由v5版本的IPluginV2Ext，到v6版本的IPluginV2IOExt和IPluginV2DynamicExt。

添加自定义TensorRT Plugin的方式有多种，以TensorRT 7.1.3为例，我平常主要用到以下两种方式：

方式一：在官方的plugin库里面进行修改添加，然后重新编译官方的plugin库，  这会生成新的一系列库文件，包括：libnvcaffeparser.so.7.1.3 ,libnvinfer_plugin.so.7.1.3，libnvonnxparser.so.7.1.3等。在TensorRT推理项目中链接这些新的动态库即可。

方式二：将自定义Plugin代码直接加入到推理项目中一同编译。

无论使用哪种方式，TensorRT Plugin代码实现的规则是一致的，我这里使用第二种方式。

实现自定义NMS3D TensorRT Plugin

        TensorRT NMS3D Plugin的实现大概有这么3个部分：

Part1：Kernel代码实现，这部分是该layer需要做的具体的CUDA操作，一般放在xxx.cu、xxx.h文件中;
Part2：定义NMS3DPlugin，继承自IPluginV2DynamicExt，是插件类，实现Plugin的主要功能;
Part3：定义NMS3DIPluginCreator，该类用于在build network时创建plugin对象，或者inference时deserialize创建plugin对象。

准备文件

我准备了以下4个文件：nms3d.cpp,nms3d_kernels.cu,nms3d.h和nms3d_kernels.h。

NMS3DPlugin插件类

class NMS3DPlugin : public nvinfer1::IPluginV2DynamicExt
{
public:
    NMS3D(const std::vector<float>& anchor_sizes, const std::vector<float>& anchor_bottom_heights);
    NMS3D(const void* data, size_t length);
    NMS3D() = delete;
    ~NMS3D() override;
    int getNbOutputs() const override;
    DimsExprs getOutputDimensions(int index, const DimsExprs* inputs, int nbInputDims, 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 _nms_post_maxsize = 100;
    int _nms_pre_maxsize = 4000 ;
    float _nms_thresh = 0.01;
    float _score_thresh = 0.3;
    int _bs = 1;
    std::vector<float> _anchor_sizes;
    std::vector<float> _anchor_bottom_heights;
    std::string mNamespace;
};

NMS3DIPluginCreator类

class NMS3DPluginCreator : public nvinfer1::IPluginCreator
{
public: 
    NMS3DPluginCreator();
    ~NMS3DPluginCreator() override = default;
    const char* getPluginName() const override;
    const char* getPluginVersion() const override;
    const PluginFieldCollection* getFieldNames() override;
    IPluginV2DynamicExt* createPlugin(const char* name, const PluginFieldCollection* fc) override;
    IPluginV2DynamicExt* deserializePlugin(const char* name, const void* serialData, size_t serialLength) override;
    void setPluginNamespace(const char* pluginNamespace) noexcept override;
    const char* getPluginNamespace() const noexcept override;
private:
    static PluginFieldCollection mFC;
    static std::vector<PluginField> mPluginAttributes;
    std::string mNamespace;
};

Plugin的注册流程

        注册插件有两种方式，第一种可以参考官方的plugin代码，在plugin/inferPlugin.cpp的initLibNvInferPlugins函数种加入自定义的Plugin。

template <typename CreatorType>
void initializePlugin(void* logger, const char* libNamespace)
{
    PluginCreatorRegistry::getInstance().addPluginCreator<CreatorType>(logger, libNamespace);
}

} // namespace plugin
} // namespace nvinfer1
// New Plugin APIs

extern "C" {
bool initLibNvInferPlugins(void* logger, const char* libNamespace)
{
    initializePlugin<nvinfer1::plugin::BatchTilePluginCreator>(logger, libNamespace);
    initializePlugin<nvinfer1::plugin::BatchedNMSPluginCreator>(logger, libNamespace);
    initializePlugin<nvinfer1::plugin::BatchedNMSDynamicPluginCreator>(logger, libNamespace);
    initializePlugin<nvinfer1::plugin::CoordConvACPluginCreator>(logger, libNamespace);
    ........
    initializePlugin<nvinfer1::plugin::NMS3DPluginCreator>(logger, libNamespace);
    return true;
}

第二种方式，可以通过REGISTER_TENSORRT_PLUGIN直接注册。我使用REGISTER_TENSORRT_PLUGIN的方式注册NMS3DPlugin。

REGISTER_TENSORRT_PLUGIN(NMS3DPluginCreator);

【参考文献】

​​​​​​TensorRT NMS的插件实现 - 知乎

​​​​​​官方文档_火鸟工作室的博客-CSDN博客