实现TensorRT-7.0插件自由!(如果不踩坑使用TensorRT插件功能)

本系列为新TensorRT的第一篇,为什么叫新,因为之前已经写了两篇关于TensorRT的文章,是关于TensorRT-5.0版本的。好久没写关于TensorRT的文章了,所幸就以来开头吧~

接下来将要讲解的TensorRT,将会是基于7.0版本。

7版本开头的TensorRT变化还是挺大的,增加了很多新特性,但是TensorRT的核心运作方式还是没有什么变化的,关于TensorRT的介绍可以看之前写的这两篇:

利用TensorRT对深度学习进行加速

利用TensorRT实现神经网络提速(读取ONNX模型并运行)

本文的内容呢,主要是讲解:

TensorRT自定义插件的使用方式

如何添加自己的自定义算子

看完本篇可以让你少踩巨多坑,客官记得常来看啊。

前言

随着tensorRT的不断发展(v5->v6->v7),TensorRT的插件的使用方式也在不断更新。插件接口也在不断地变化,由v5版本的IPluginV2Ext,到v6版本的IPluginV2IOExt和IPluginV2DynamicExt。未来不知道会不会出来新的API,不过这也不是咱要考虑的问题,因为TensorRT的后兼容性做的很好,根本不用担心你写的旧版本插件在新版本上无法运行。

目前的plugin-API:

QQ20201103-101737

TensorRT插件的存在目的,主要是为了让我们实现TensorRT目前还不支持的算子,毕竟众口难调嘛,我们在转换过程中肯定会有op不支持的情况。这个时候就需要使用TensorRT的plugin去实现我们的自己的op。此时我们需要通过TensorRT提供的接口去实现自己的op,因此这个plugin的生命周期也需要遵循TensorRT的规则

一个简单的了解

那么plugin到底长啥样,可以先看看TensorRT的官方plugin库长啥样,截止写这篇文章时,master分支是7.2版本的plugin:

https://github.com/NVIDIA/TensorRT/tree/master/plugin

tensorrt-plugin

官方提供的插件已经相当多,而且TensorRT开源了plugin部分(可以让我们白嫖!)。并且可以看到其源码,通过模仿源码来学习plugin是如何写的。

如果要添加自己的算子,可以在官方的plugin库里头进行修改添加,然后编译官方的plugin库。将生成的libnvinfer_plugin.so.7替换原本的.so文件即可。或者自己写一个类似于官方plugin的组件,将名称替换一下,同样生成.so,在TensorRT的推理项目中引用这个动态链接库即可。

以下介绍中,我们需要写的IPlugin简称为插件op。

开始写插件

有兴趣的可以先看看TensorRT的,官方文档的介绍简单意骇,不过坑是少不了的..而本文的目的,就是尽量让你少趟坑。

首先按照官方plugin的排布方式,下面随便挑了个官方plugin:

instance_normalization_plugin

准备一个自己的插件:custom.cpp和custom.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类型:

内容版权声明:除非注明,否则皆为本站原创文章。

转载注明出处:https://www.heiqu.com/wpwjpf.html