TensorRT5.1.5.0入門 自定義層 IPlugin & IPluginV2的對比(C++)

最近研究 TensorRT的自定義層，嘗試的使用了自定義的FC層FC層和Upsample層Upsample層之後，重新回去看開發者手冊，在此記錄。
　　自己的理解，TensorRT的自定義層機制是有兩個方法的，一種基於基類IPlugin，另一種是基於基類IPluginV2，從字面意思上來看IPluginV2就是最新版本。
　　IPlugin類的方法，是通過自己編寫IPlugin的派生類IPluginExt和nvinfer1::IPluginFactory類以及nvcaffeparser1::IPluginFactoryExt類來實現自定義層的編輯和使用的。具體的實例可以看上面的鏈接。就是通過IPluginFactory類將IPluginExt類定義的Plugin實例化，然後在TensorRT中使用。在Developer Guide中對應的是如下這段代碼：

//The following sample code adds a new plugin called FooPlugin:
class FooPlugin : public IPluginExt
{
	...implement all class methods for your plugin
};

class MyPluginFactory : public nvinfer1::IPluginFactory, public nvcaffeparser1::IPluginFactoryExt
{
	...implement all factory methods for your plugin
};

.
　　IPluginV2類的方法，是關於自定義層的另一種方法，這種方法是調用已經被註冊的plugin的方法，這種方法不依賴nvinfer1::IPluginFactory去實例化，而是使用nvcaffeparser1::IPluginFactoryV2 和 IPluginCreator代替，這種方法對應Developer Guide中的如下代碼：

class FooPlugin : public IPluginV2
{
	...implement all class methods for your plugin
};

class FooPluginFactory : public nvcaffeparser1::IPluginFactoryV2
{
	virtual nvinfer1::IPluginV2* createPlugin(...)
	{
		...create and return plugin object of type FooPlugin
	}
	bool isPlugin(const char* name) 
	{
		...check if layer name corresponds to plugin
	}
}

class FooPluginCreator : public IPluginCreator
{
	...implement all creator methods here
};
REGISTER_TENSORRT_PLUGIN(FooPluginCreator);

以前的博客已經用過了很多次基於IPlugin類的自定義層構造，我稱之爲V1用法，現在主要介紹兩者的對比以及IPluginV2類的自定義層用法，稱之爲V2用法。

IPlugin類和IPluginV2類的區別

IPlugin類 & IPluginV2類在NvInfer.h中的聲明

class IPlugin
{
public:
	virtual int getNbOutputs() const = 0;
	virtual Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) = 0;
	virtual void configure(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, int maxBatchSize) = 0;
	virtual int initialize() = 0;
	virtual size_t getWorkspaceSize(int maxBatchSize) const = 0;
	virtual int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) = 0;
	virtual size_t getSerializationSize() = 0;
	virtual void serialize(void* buffer) = 0;
	virtual ~IPlugin() {}
};
class IPluginExt
{
public:
	virtual int getTensorRTVersion() const
	virtual bool supportsFormat(DataType type, PluginFormat format) const = 0;
	virtual void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) = 0;
	virtual ~IPluginExt() {}
}

class IPluginV2
{
public:
	virtual int getTensorRTVersion() const { return NV_TENSORRT_VERSION;}
	virtual const char* getPluginType() const = 0;
	virtual const char* getPluginVersion() const = 0;
	virtual int getNbOutputs() const = 0;
	virtual Dims getOutputDimensions(int index, const Dims* inputs, int nbInputDims) = 0;
	virtual bool supportsFormat(DataType type, PluginFormat format) const = 0;
	virtual void configureWithFormat(const Dims* inputDims, int nbInputs, const Dims* outputDims, int nbOutputs, DataType type, PluginFormat format, int maxBatchSize) = 0;
	virtual int initialize() = 0; 
	virtual void terminate() = 0;
	virtual size_t getWorkspaceSize(int maxBatchSize) const = 0;
	virtual int enqueue(int batchSize, const void* const* inputs, void** outputs, void* workspace, cudaStream_t stream) = 0;
	virtual size_t getSerializationSize() const = 0;
	virtual void serialize(void* buffer) const = 0;
	virtual void destroy() = 0;
	virtual IPluginV2* clone() const = 0;
	virtual void setPluginNamespace(const char* pluginNamespace) = 0;
	virtual const char* getPluginNamespace() const = 0;
protected:
    virtual ~IPluginV2() {}
};

.
　　從類聲明的對比來看，IPlugin類比IPluginV2類少了很多方法，但是IPlugin方法最後new出來的Plugin是繼承於IPluginExt，IPluginExt才又繼承與IPlugin，V1方法在IPluginExt類中聲明瞭getTensorRTVersion()，supportsFormat(), configureWithFormat(),configure()等方法。
　　
　　所以，比較起來IPluginV2多了getPluginVersion(), getPluginType(), setPluginNamespace(), getPluginNamespace()和IPluginV2* clone()方法。
　　getPluginType()是用來匹配 plugin creator 返回的plugin name的方法。
　　getPluginVersion()是用來匹配 plugin creator 返回的plugin version的方法。比如在Developer Guide中，提到的這些RPROI_TRT， Normalize_TRT， PriorBox_TRT等已經被tensorRT封裝好的Plugin都是屬於version1的。
　　setPluginNamespace()是用來設置這個plugin對象屬於哪個namesapce的方法，在相同plugin library中的所有plugin對象都應該在同一個namespace中。
　　getPluginNamespace()是用來返回該plugin對象所屬namespace的。
　　IPluginV2* clone()方法，每次創建包含此plugin的新builder，network或engine時，都會調用此方法。 它應該返回一個帶有正確參數的新plugin對象。
　　其餘的方法都基本相同，說明其實IPluginV2類和IPlugin類的定義方法相似。

IPluginV2Ext類：擴展IPluginV2的功能

但是，其實除了上面列出的和IPlugin類類似的方法之外，IPluginV2還有一個派生類IPluginExtV2，通過支持不同的output數據類型和broadcast across batch擴展了IPluginV2類的功能。

class IPluginV2Ext : public IPluginV2
{
public:
	virtual nvinfer1::DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const = 0;
	virtual bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const = 0;
	virtual bool canBroadcastInputAcrossBatch(int inputIndex) const = 0;
	virtual void configurePlugin(const Dims* inputDims, int nbInputs, const Dims* outputDims,
                                 int nbOutputs, const DataType* inputTypes, const DataType* outputTypes,
                                 const bool* inputIsBroadcast, const bool* outputIsBroadcast, PluginFormat floatFormat, int maxBatchSize)=0;
     virtual ~IPluginV2Ext() {}
     virtual void attachToContext(cudnnContext* /*cudnn*/, cublasContext* /*cublas*/, IGpuAllocator* /*allocator*/) {}
     virtual void detachFromContext() {}
     virtual IPluginV2Ext* clone() const _TENSORRT_OVERRIDE = 0;
protected:
	int getTensorRTVersion() const _TENSORRT_OVERRIDE
    {
        return (0x01000000 | (NV_TENSORRT_VERSION & 0xFFFFFF));
    }
    void configureWithFormat(const Dims* /*inputDims*/, int /*nbInputs*/, const Dims* /*outputDims*/,
                             int /*nbOutputs*/, DataType /*type*/, PluginFormat /*format*/, int /*maxBatchSize*/) _TENSORRT_OVERRIDE _TENSORRT_FINAL {}
};                          

.
　　這裏挖個坑，這些方法的具體代碼實踐以後再試試看。

IPluginCreator類：實現自定義層的創建

除了類中方法和派生類的不同，IPluginV2的自定義層的創建，還利用到了另一個類：IPluginCreator。
這個類是用來在Plugin Registry中查找和創建相應的plugin。

class IPluginCreator
{
public:
	virtual int getTensorRTVersion() const { return NV_TENSORRT_VERSION; }
	virtual const char* getPluginName() const = 0;
	virtual const char* getPluginVersion() const = 0;
	virtual const PluginFieldCollection* getFieldNames() = 0;
	virtual IPluginV2* createPlugin(const char* name, const PluginFieldCollection* fc) = 0;
	virtual IPluginV2* deserializePlugin(const char* name, const void* serialData, size_t serialLength) = 0;
	virtual void setPluginNamespace(const char* pluginNamespace) = 0;
	virtual const char* getPluginNamespace() const = 0;
	virtual ~IPluginCreator() {}
};

.
　　getPluginName()方法用來返回plugin name 和匹配IPluginExt::getPluginType的返回值
　　getPluginVersion()方法返回plugin version，默認是1
　　getFieldNames()方法返回PluginFieldCollection結構，這個結構會依據field name去填充相關字段，填充PluginFieldType等。
　　createPlugin()方法就是用PluginFieldCollection來創建Plugin，對應的數據會被填充。
　　deserializePlugin()方法返回用於推理的Plugin對象，TensorRT根據plugin名稱和版本在內部調用這個對象。
　　set/getPluginNamespace()方法用於設置或者獲取這個Creator實例所屬的namespace

基於IPluginV2類的自定義層的用法

這個部分準備好好研究一下，所以單寫一個博客。
我的想法是分爲
（1）IPluginV2方法使用tensorRT已有的自定義層
（2）IPluginV2方法使用自己定義的自定義層
具體請看這裏IPluginV2