refactor: embedding_head

* add abstract poolinglayer

* add poolinglayer factory

projects/FastRT/fastrt/factory/CMakeLists.txt

@@ -1,4 +1,5 @@
 target_sources(${PROJECT_NAME}
 PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/factory.cpp
+  ${CMAKE_SOURCE_DIR}/fastrt/layers/poolingLayerRT.h
 )

projects/FastRT/fastrt/factory/factory.cpp

@@ -3,6 +3,7 @@
 #include "fastrt/sbs_resnet.h"
 #include "fastrt/factory.h"
 #include "fastrt/embedding_head.h"
+#include "../layers/poolingLayerRT.h"
 
 namespace fastrt {
 
@@ -46,4 +47,24 @@ namespace fastrt {
         }
     }
 
+    std::unique_ptr<IPoolingLayerRT> LayerFactory::createPoolingLayer(const FastreidPoolingType& pooltype) {
+        switch(pooltype) {
+            case FastreidPoolingType::maxpool:
+                std::cout << "[createPoolingLayer]: maxpool" << std::endl;
+                return make_unique<MaxPool>();
+            case FastreidPoolingType::avgpool:
+                std::cout << "[createPoolingLayer]: avgpool" << std::endl;
+                return make_unique<AvgPool>();
+            case FastreidPoolingType::gempool:
+                std::cout << "[createPoolingLayer]: gempool" << std::endl;
+                return make_unique<GemPool>();
+            case FastreidPoolingType::gempoolP:
+                std::cout << "[createPoolingLayer]: gempoolP" << std::endl;
+                return make_unique<GemPoolP>();
+            default:
+                std::cerr << "[Pooling layer is not supported.]" << std::endl; 
+                return nullptr;
+        }  
+    }
+
 }

projects/FastRT/fastrt/heads/embedding_head.cpp

@@ -5,37 +5,16 @@
 
 namespace fastrt {
 
+    embedding_head::embedding_head() : _layerFactory(make_unique<LayerFactory>()) {}
+    
+    embedding_head::embedding_head(std::unique_ptr<LayerFactory> layerFactory) : _layerFactory(std::move(layerFactory)) {}
+
     ILayer* embedding_head::topology(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
         /*
          * Reference: https://github.com/JDAI-CV/fast-reid/blob/master/fastreid/modeling/heads/embedding_head.py
          */
-        ILayer* pooling{nullptr};
-        switch(reidCfg.pooling) {
-            case FastreidPoolingType::maxpool:
-                pooling = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});
-                { 
-                    auto p = dynamic_cast<IPoolingLayer*>(pooling);
-                    if(p) p->setStrideNd(DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});
-                    else std::cout << "Downcasting failed." << std::endl; 
-                }
-                break;
-            case FastreidPoolingType::avgpool:
-                pooling = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});
-                { 
-                    auto p = dynamic_cast<IPoolingLayer*>(pooling);
-                    if(p) p->setStrideNd(DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});
-                    else std::cout << "Downcasting failed." << std::endl; 
-                }
-                break;
-            case FastreidPoolingType::gempool:
-                pooling = trtxapi::addGeneralizedMeanPooling(network, input); 
-                break;
-            case FastreidPoolingType::gempoolP:
-                pooling = trtxapi::addGeneralizedMeanPooling(network, input, *(float*)weightMap["heads.pool_layer.p"].values); 
-                break;
-            default:
-                std::cout << "This pooling layer is not supported." << std::endl; 
-        }
+
+        ILayer* pooling = _layerFactory->createPoolingLayer(reidCfg.pooling)->addPooling(network, weightMap, input);
         TRTASSERT(pooling);
 
         // Hint: It's used to be "heads.bnneck.0" before Sep 10, 2020. (JDAI-CV/fast-reid)
@@ -57,5 +36,5 @@ namespace fastrt {
         TRTASSERT(bottleneck);
         return bottleneck;
     }
-    
+
 }

projects/FastRT/fastrt/layers/CMakeLists.txt

@@ -1,4 +1,6 @@
 target_sources(${PROJECT_NAME}
 PRIVATE
   ${CMAKE_CURRENT_SOURCE_DIR}/layers.cpp
+  ${CMAKE_CURRENT_SOURCE_DIR}/poolingLayerRT.h
+  ${CMAKE_CURRENT_SOURCE_DIR}/poolingLayerRT.cpp
 )

projects/FastRT/fastrt/layers/poolingLayerRT.cpp

@@ -0,0 +1,31 @@
+#include <iostream>
+#include "fastrt/layers.h"
+#include "poolingLayerRT.h"
+
+namespace fastrt {
+
+    ILayer* MaxPool::addPooling(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input) {
+        ILayer* pooling = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});       
+        auto p = dynamic_cast<nvinfer1::IPoolingLayer*>(pooling);
+        if(p) p->setStrideNd(DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});
+        else std::cout << "Downcasting failed." << std::endl; 
+        return pooling;
+    }
+    
+    ILayer* AvgPool::addPooling(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input) {
+        ILayer* pooling = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});
+        auto p = dynamic_cast<IPoolingLayer*>(pooling);
+        if(p) p->setStrideNd(DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]});
+        else std::cout << "Downcasting failed." << std::endl; 
+        return pooling;
+    }
+
+    ILayer* GemPool::addPooling(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input) {
+        return trtxapi::addGeneralizedMeanPooling(network, input); 
+    }
+
+    ILayer* GemPoolP::addPooling(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input) {
+        return trtxapi::addGeneralizedMeanPooling(network, input, *(float*)weightMap["heads.pool_layer.p"].values); 
+    }    
+
+}

projects/FastRT/fastrt/layers/poolingLayerRT.h

@@ -0,0 +1,46 @@
+#include "NvInfer.h"
+#include "fastrt/IPoolingLayerRT.h"
+using namespace nvinfer1;
+
+namespace fastrt {
+
+    class MaxPool : public IPoolingLayerRT {
+    public:
+        MaxPool() = default;
+        ~MaxPool() = default;
+
+        ILayer* addPooling(INetworkDefinition *network, 
+            std::map<std::string, Weights>& weightMap,
+            ITensor& input) override;
+    };
+
+    class AvgPool : public IPoolingLayerRT {
+    public:
+        AvgPool() = default;
+        ~AvgPool() = default;
+
+        ILayer* addPooling(INetworkDefinition *network, 
+            std::map<std::string, Weights>& weightMap,
+            ITensor& input) override;
+    };
+
+    class GemPool : public IPoolingLayerRT {
+    public:
+        GemPool() = default;
+        ~GemPool() = default;
+
+        ILayer* addPooling(INetworkDefinition *network, 
+            std::map<std::string, Weights>& weightMap,
+            ITensor& input) override;
+    };
+
+    class GemPoolP : public IPoolingLayerRT {
+    public:
+        GemPoolP() = default;
+        ~GemPoolP() = default;
+
+        ILayer* addPooling(INetworkDefinition *network, 
+            std::map<std::string, Weights>& weightMap,
+            ITensor& input) override;
+    };
+}

projects/FastRT/include/fastrt/IPoolingLayerRT.h

@@ -0,0 +1,20 @@
+#pragma once
+
+#include <map>
+#include "struct.h"
+#include "NvInfer.h"
+using namespace nvinfer1;
+
+namespace fastrt {
+
+    class IPoolingLayerRT {
+    public:
+        IPoolingLayerRT() = default;
+        virtual ~IPoolingLayerRT() = default;
+
+        virtual ILayer* addPooling(INetworkDefinition *network, 
+            std::map<std::string, Weights>& weightMap, 
+            ITensor& input) = 0; 
+    };
+
+}

projects/FastRT/include/fastrt/embedding_head.h

@@ -4,13 +4,18 @@
 #include "NvInfer.h"
 #include "fastrt/module.h"
 #include "fastrt/struct.h"
+#include "fastrt/factory.h"
 using namespace nvinfer1;
 
 namespace fastrt {
 
     class embedding_head : public Module {
+    private:
+        std::unique_ptr<LayerFactory> _layerFactory;
+
     public:
-        embedding_head() = default;
+        embedding_head();
+        embedding_head(std::unique_ptr<LayerFactory> layerFactory);
         ~embedding_head() = default;
 
         ILayer* topology(INetworkDefinition *network, 
@@ -18,4 +23,5 @@ namespace fastrt {
             ITensor& input, 
             const FastreidConfig& reidCfg) override;
     };
+
 }

projects/FastRT/include/fastrt/factory.h

@@ -2,6 +2,7 @@
 
 #include "struct.h"
 #include "module.h"
+#include "IPoolingLayerRT.h"
 
 namespace fastrt {
     
@@ -14,4 +15,12 @@ namespace fastrt {
         std::unique_ptr<Module> createHead(const FastreidHeadType& headtype);
     };
 
+    class LayerFactory {
+    public:
+        LayerFactory() = default;
+        ~LayerFactory() = default;
+
+        std::unique_ptr<IPoolingLayerRT> createPoolingLayer(const FastreidPoolingType& pooltype);
+    };
+
 }