TensorRT学习笔记6 - IPlugin

简介

       IPlugin是一个类，用于实现自定义层（custom layer），我们可以继承这个类，然后实现它的成员函数。在使用IPlugin时，要include

构造函数

       继承IPlugin类需要实现两个构造函数。

• 构造函数1
         plugin(const int parameter1, const int parameter2, …)
         此构造函数和我们常写的构造函数，没有什么区别。
• 构造函数2
         plugin(const void* data, size_t length)
         此构造函数通过反序列化进行参数读入，和成员函数serialize对应。

析构函数

       析构函数就是我们常写的析构函数，没有什么区别。

成员函数

• configure
         virtual void nvinfer1::IPlugin::configure(
                const Dims * inputDims, -----------------------------> 输入tensor的维度
                int nbInputs, ----------------------------------------------> 输入个数
                const Dims * outputDims, ----------------------------> 输出tensor的维度
                int nbOutputs, --------------------------------------------> 输出个数
                int maxBatchSize ---------------------------------------> 最大的BatchSize)
         函数功能：配置自定义层，这个函数会在调用initialize()之前调用
         Returns：没有返回值
• enqueue
         virtual int nvinfer1::IPlugin::enqueue(
                int batchSize, ---------------------------------------------> BatchSize
                const void *const * inputs, ----------------------------> 输入tensor的内存
                void ** outputs, -------------------------------------------> 输出tensor的内存
                void * workspace, ---------------------------------------> 执行的工作空间
                cudaStream_t stream ----------------------------------> 执行kernel函数的stream)
         函数功能：执行自定义层，具体的过程一般写在cu文件中
         Returns：返回0代表success，返回其他值会终止执行。
• getNbOutputs
         virtual int nvinfer1::IPlugin::getNbOutputs( ) const
         函数功能：获得自定义层的输出个数
         Returns：自定义层的输出个数
• getOutputDimensions
         virtual Dims nvinfer1::IPlugin::getOutputDimensions(
                int index, ---------------------------------------------------> 输出tensor的index
                const Dims * inputs, ------------------------------------> 输入tensor
                int nbInputDims ------------------------------------------> 输入个数)
         函数功能：获得自定义层的输出tensor的维度信息
         Returns：自定义层的输出的维度信息
• getSerializationSize
         virtual size_t nvinfer1::IPlugin::getSerializationSize( )
         函数功能：计算序列化缓冲区需要的size
         Returns： 序列化缓冲区的size
• getWorkspaceSize
         virtual size_t nvinfer1::IPlugin::getWorkspaceSize(
                int maxBatchSize ---------------------------------------> 最大的BatchSize) const
         函数功能：计算工作区间需要的size
         Returns：工作区间的size
• initialize
         virtual int nvinfer1::IPlugin::initialize( )
         函数功能：在执行前，初始化自定义层，此函数会在创建engine时调用
         Returns：返回0代表success，返回其他值会终止执行。
• serialize
         virtual void nvinfer1::IPlugin::serialize(
                void * buffer -----------------------------------------------> 一个指向size>getSerializationSize()的缓冲区的指针)
         函数功能：将自定义层序列化
         Returns：没有返回值
• terminate
         virtual void nvinfer1::IPlugin::terminate( )
         函数功能：释放自定义层初始化时请求的资源，此函数会在销毁engine时调用
         Returns：没有返回值