From 0617e3eeb720dfd267333e2a5149260eac7712ac Mon Sep 17 00:00:00 2001 From: darrenhsieh Date: Sat, 13 Feb 2021 21:28:29 +0800 Subject: [PATCH] refactor: embedding_head * add abstract poolinglayer * add poolinglayer factory --- projects/FastRT/fastrt/factory/CMakeLists.txt | 1 + projects/FastRT/fastrt/factory/factory.cpp | 21 +++++++++ .../FastRT/fastrt/heads/embedding_head.cpp | 35 +++----------- projects/FastRT/fastrt/layers/CMakeLists.txt | 2 + .../FastRT/fastrt/layers/poolingLayerRT.cpp | 31 +++++++++++++ .../FastRT/fastrt/layers/poolingLayerRT.h | 46 +++++++++++++++++++ .../FastRT/include/fastrt/IPoolingLayerRT.h | 20 ++++++++ .../FastRT/include/fastrt/embedding_head.h | 8 +++- projects/FastRT/include/fastrt/factory.h | 9 ++++ 9 files changed, 144 insertions(+), 29 deletions(-) create mode 100644 projects/FastRT/fastrt/layers/poolingLayerRT.cpp create mode 100644 projects/FastRT/fastrt/layers/poolingLayerRT.h create mode 100644 projects/FastRT/include/fastrt/IPoolingLayerRT.h diff --git a/projects/FastRT/fastrt/factory/CMakeLists.txt b/projects/FastRT/fastrt/factory/CMakeLists.txt index b1797fe..7b4e833 100644 --- a/projects/FastRT/fastrt/factory/CMakeLists.txt +++ b/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 ) \ No newline at end of file diff --git a/projects/FastRT/fastrt/factory/factory.cpp b/projects/FastRT/fastrt/factory/factory.cpp index c1d1c5b..c2ae092 100644 --- a/projects/FastRT/fastrt/factory/factory.cpp +++ b/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 LayerFactory::createPoolingLayer(const FastreidPoolingType& pooltype) { + switch(pooltype) { + case FastreidPoolingType::maxpool: + std::cout << "[createPoolingLayer]: maxpool" << std::endl; + return make_unique(); + case FastreidPoolingType::avgpool: + std::cout << "[createPoolingLayer]: avgpool" << std::endl; + return make_unique(); + case FastreidPoolingType::gempool: + std::cout << "[createPoolingLayer]: gempool" << std::endl; + return make_unique(); + case FastreidPoolingType::gempoolP: + std::cout << "[createPoolingLayer]: gempoolP" << std::endl; + return make_unique(); + default: + std::cerr << "[Pooling layer is not supported.]" << std::endl; + return nullptr; + } + } + } \ No newline at end of file diff --git a/projects/FastRT/fastrt/heads/embedding_head.cpp b/projects/FastRT/fastrt/heads/embedding_head.cpp index 8209606..6d0dd8f 100644 --- a/projects/FastRT/fastrt/heads/embedding_head.cpp +++ b/projects/FastRT/fastrt/heads/embedding_head.cpp @@ -5,37 +5,16 @@ namespace fastrt { + embedding_head::embedding_head() : _layerFactory(make_unique()) {} + + embedding_head::embedding_head(std::unique_ptr layerFactory) : _layerFactory(std::move(layerFactory)) {} + ILayer* embedding_head::topology(INetworkDefinition *network, std::map& 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(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(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; } - + } \ No newline at end of file diff --git a/projects/FastRT/fastrt/layers/CMakeLists.txt b/projects/FastRT/fastrt/layers/CMakeLists.txt index 1ba966e..7d565ab 100644 --- a/projects/FastRT/fastrt/layers/CMakeLists.txt +++ b/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 ) \ No newline at end of file diff --git a/projects/FastRT/fastrt/layers/poolingLayerRT.cpp b/projects/FastRT/fastrt/layers/poolingLayerRT.cpp new file mode 100644 index 0000000..884a0b3 --- /dev/null +++ b/projects/FastRT/fastrt/layers/poolingLayerRT.cpp @@ -0,0 +1,31 @@ +#include +#include "fastrt/layers.h" +#include "poolingLayerRT.h" + +namespace fastrt { + + ILayer* MaxPool::addPooling(INetworkDefinition *network, std::map& weightMap, ITensor& input) { + ILayer* pooling = network->addPoolingNd(input, PoolingType::kMAX, DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]}); + auto p = dynamic_cast(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& weightMap, ITensor& input) { + ILayer* pooling = network->addPoolingNd(input, PoolingType::kAVERAGE, DimsHW{input.getDimensions().d[1], input.getDimensions().d[2]}); + auto p = dynamic_cast(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& weightMap, ITensor& input) { + return trtxapi::addGeneralizedMeanPooling(network, input); + } + + ILayer* GemPoolP::addPooling(INetworkDefinition *network, std::map& weightMap, ITensor& input) { + return trtxapi::addGeneralizedMeanPooling(network, input, *(float*)weightMap["heads.pool_layer.p"].values); + } + +} \ No newline at end of file diff --git a/projects/FastRT/fastrt/layers/poolingLayerRT.h b/projects/FastRT/fastrt/layers/poolingLayerRT.h new file mode 100644 index 0000000..b74bcef --- /dev/null +++ b/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& weightMap, + ITensor& input) override; + }; + + class AvgPool : public IPoolingLayerRT { + public: + AvgPool() = default; + ~AvgPool() = default; + + ILayer* addPooling(INetworkDefinition *network, + std::map& weightMap, + ITensor& input) override; + }; + + class GemPool : public IPoolingLayerRT { + public: + GemPool() = default; + ~GemPool() = default; + + ILayer* addPooling(INetworkDefinition *network, + std::map& weightMap, + ITensor& input) override; + }; + + class GemPoolP : public IPoolingLayerRT { + public: + GemPoolP() = default; + ~GemPoolP() = default; + + ILayer* addPooling(INetworkDefinition *network, + std::map& weightMap, + ITensor& input) override; + }; +} \ No newline at end of file diff --git a/projects/FastRT/include/fastrt/IPoolingLayerRT.h b/projects/FastRT/include/fastrt/IPoolingLayerRT.h new file mode 100644 index 0000000..a5a017d --- /dev/null +++ b/projects/FastRT/include/fastrt/IPoolingLayerRT.h @@ -0,0 +1,20 @@ +#pragma once + +#include +#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& weightMap, + ITensor& input) = 0; + }; + +} \ No newline at end of file diff --git a/projects/FastRT/include/fastrt/embedding_head.h b/projects/FastRT/include/fastrt/embedding_head.h index 6de8fb3..32c8314 100644 --- a/projects/FastRT/include/fastrt/embedding_head.h +++ b/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; + public: - embedding_head() = default; + embedding_head(); + embedding_head(std::unique_ptr layerFactory); ~embedding_head() = default; ILayer* topology(INetworkDefinition *network, @@ -18,4 +23,5 @@ namespace fastrt { ITensor& input, const FastreidConfig& reidCfg) override; }; + } \ No newline at end of file diff --git a/projects/FastRT/include/fastrt/factory.h b/projects/FastRT/include/fastrt/factory.h index 7d17432..80fd674 100644 --- a/projects/FastRT/include/fastrt/factory.h +++ b/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 createHead(const FastreidHeadType& headtype); }; + class LayerFactory { + public: + LayerFactory() = default; + ~LayerFactory() = default; + + std::unique_ptr createPoolingLayer(const FastreidPoolingType& pooltype); + }; + } \ No newline at end of file