(tkDNN): Support Resnet-101-AP-GeM for CUDNN, TRT to be fixed

Signed-off-by: Micaela Verucchi <micaelaverucchi@gmail.com>
2022-09-24 16:57:12 +02:00
parent d4f7b4ad8b
commit 56da10df64
9 changed files with 741 additions and 3 deletions
@@ -488,7 +488,8 @@ typedef enum {
    POOLING_MAX     = 0,
    POOLING_AVERAGE = 1,                    // count for average includes padded values
    POOLING_AVERAGE_EXCLUDE_PADDING = 2,    // count for average does not include padded values
-    POOLING_MAX_FIXEDSIZE = 100             // max pool darknet fashion
+    POOLING_MAX_FIXEDSIZE = 100,             // max pool darknet fashion
+    POOLING_GENERALIZED_MEAN_P = 200       // mean pooling with pow parameter 
 } tkdnnPoolingMode_t;

 /**
@@ -498,6 +499,7 @@ typedef enum {
 class Pooling : public Layer {

 public:
+    float pow_param;
    int winH, winW;
    int strideH, strideW;
    int paddingH, paddingW;
@@ -508,7 +510,7 @@ public:
    Pooling(Network *net, int winH, int winW, 
            int strideH, int strideW, 
            int paddingH, int paddingW,
-            tkdnnPoolingMode_t pool_mode);
+            tkdnnPoolingMode_t pool_mode, float p = 1.0f);
    virtual ~Pooling();
    virtual layerType_t getLayerType() { return LAYER_POOLING; };

@@ -20,6 +20,8 @@ void reorgForward(dnnType *srcData, dnnType *dstData,

 void MaxPoolingForward(dnnType *srcData, dnnType *dstData, int n, int c, int h, int w, int stride_x, int stride_y, int size, int padding, cudaStream_t stream = cudaStream_t(0));

+void GeneralizedMeanPoolingP(dnnType* srcData, dnnType* dstData, int n, int c, int h, int w, float p, cudaStream_t stream  = cudaStream_t(0));
+
 void softmaxForward(float *input, int n, int batch, int batch_offset,
                    int groups, int group_offset, int stride, float temp, float *output, cudaStream_t stream = cudaStream_t(0));

@@ -0,0 +1,103 @@
+#include<cassert>
+#include "../kernels.h"
+#include <NvInfer.h>
+#include <vector>
+#include <utils.h>
+
+
+namespace nvinfer1 {
+    class GeneralizedMeanPoolingPRT : public IPluginV2Ext {
+
+    public:
+        GeneralizedMeanPoolingPRT(int input_c, int input_h, int input_w, int input_n, int output_c, int output_h, int output_w, int output_n, float p) ;
+
+        GeneralizedMeanPoolingPRT(const void *data, size_t length) ;
+
+        ~GeneralizedMeanPoolingPRT() ;
+
+        int getNbOutputs() const NOEXCEPT override ;
+
+        Dims getOutputDimensions(int index, const Dims *inputs, int nbInputDims) NOEXCEPT override ;
+
+        int initialize() NOEXCEPT override ;
+
+        void terminate() NOEXCEPT override ;
+
+        size_t getWorkspaceSize(int maxBatchSize) const NOEXCEPT override ;
+
+#if NV_TENSORRT_MAJOR > 7
+        int enqueue(int batchSize, const void *const *inputs, void *const *outputs, void *workspace,
+                            cudaStream_t stream) NOEXCEPT override ;
+#elif NV_TENSORRT_MAJOR <= 7
+        int32_t enqueue (int32_t batchSize, const void *const *inputs, void **outputs, void *workspace, cudaStream_t stream) override;
+#endif
+
+
+        size_t getSerializationSize() const NOEXCEPT override ;
+
+        void serialize(void *buffer) const NOEXCEPT override ;
+
+        void destroy() NOEXCEPT override ;
+
+        bool supportsFormat(DataType type, PluginFormat format) const NOEXCEPT override ;
+
+        const char *getPluginNamespace() const NOEXCEPT override ;
+
+        void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
+
+        const char *getPluginType() const NOEXCEPT override ;
+
+        const char *getPluginVersion() const NOEXCEPT override ;
+
+        IPluginV2Ext *clone() const NOEXCEPT override ;
+
+        DataType getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const NOEXCEPT override;
+
+        void attachToContext(cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) NOEXCEPT override;
+
+        bool isOutputBroadcastAcrossBatch(int outputIndex, const bool* inputIsBroadcasted, int nbInputs) const NOEXCEPT override;
+
+        bool canBroadcastInputAcrossBatch(int inputIndex) const NOEXCEPT override;
+
+        void configurePlugin (Dims const *inputDims, int32_t nbInputs, Dims const *outputDims,
+                              int32_t nbOutputs, DataType const *inputTypes, DataType const *outputTypes,
+                              bool const *inputIsBroadcast, bool const *outputIsBroadcast, PluginFormat floatFormat,
+                              int32_t maxBatchSize) NOEXCEPT override;
+
+        void detachFromContext() NOEXCEPT override;
+
+
+        int i_n, i_c, i_h, i_w, o_n, o_c, o_h, o_w;
+        float p;
+
+    private:
+        std::string mPluginNamespace;
+    };
+
+    class GeneralizedMeanPoolingPRTPluginCreator : public IPluginCreator {
+    public:
+        GeneralizedMeanPoolingPRTPluginCreator() ;
+
+        void setPluginNamespace(const char *pluginNamespace) NOEXCEPT override ;
+
+        const char *getPluginNamespace() const NOEXCEPT override ;
+
+        IPluginV2Ext *deserializePlugin(const char *name, const void *serialData, size_t serialLength) NOEXCEPT override ;
+
+        IPluginV2Ext *createPlugin(const char *name, const PluginFieldCollection *fc) NOEXCEPT override ;
+
+        const char *getPluginName() const NOEXCEPT override ;
+
+        const char *getPluginVersion() const NOEXCEPT override ;
+
+        const PluginFieldCollection *getFieldNames() NOEXCEPT override ;
+
+    private:
+        static PluginFieldCollection mFC;
+        static std::vector<PluginField> mPluginAttributes;
+        std::string mPluginNamespace;
+
+    };
+
+    REGISTER_TENSORRT_PLUGIN(GeneralizedMeanPoolingPRTPluginCreator);
+};