mirrors
/
fast-reid
mirror of https://github.com/JDAI-CV/fast-reid.git
1
0
Fork 0
Code Issues Projects Releases Wiki Activity

refactor: model class 

* add Module as base class

* add head type

* replace template func args with Module ptr

* add Module factory
pull/411/head
darrenhsieh 2021-02-13 21:14:14 +08:00
parent 254a489eb1
commit 159494e4a4
17 changed files with 274 additions and 181 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
projects/FastRT/CMakeLists.txt
View File

@ -4,7 +4,7 @@ set(LIBARARY_NAME "FastRT" CACHE STRING "The Fastreid-tensorrt library name")
set(LIBARARY_VERSION_MAJOR "0")
set(LIBARARY_VERSION_MINOR "0")
set(LIBARARY_VERSION_SINOR "2")
set(LIBARARY_VERSION_SINOR "3")
set(LIBARARY_SOVERSION "0")
set(LIBARARY_VERSION "${LIBARARY_VERSION_MAJOR}.${LIBARARY_VERSION_MINOR}.${LIBARARY_VERSION_SINOR}")
project(${LIBARARY_NAME}${LIBARARY_VERSION})

12
projects/FastRT/README.md
View File

@ -84,8 +84,9 @@ static const int INPUT_W = 128;
static const int OUTPUT_SIZE = 2048;
static const int DEVICE_ID = 0;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const FastreidBackboneType BACKBONE = FastreidBackboneType::r50;
static const FastreidHeadType HEAD = FastreidHeadType::EmbeddingHead;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const int LAST_STRIDE = 1;
static const bool WITH_IBNA = true;
static const bool WITH_NL = true;
@ -103,8 +104,9 @@ static const int INPUT_W = 128;
static const int OUTPUT_SIZE = 2048;
static const int DEVICE_ID = 0;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const FastreidBackboneType BACKBONE = FastreidBackboneType::r50;
static const FastreidHeadType HEAD = FastreidHeadType::EmbeddingHead;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const int LAST_STRIDE = 1;
static const bool WITH_IBNA = false;
static const bool WITH_NL = true;
@ -122,8 +124,9 @@ static const int INPUT_W = 128;
static const int OUTPUT_SIZE = 512;
static const int DEVICE_ID = 0;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const FastreidBackboneType BACKBONE = FastreidBackboneType::r34_distill;
static const FastreidHeadType HEAD = FastreidHeadType::EmbeddingHead;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const int LAST_STRIDE = 1;
static const bool WITH_IBNA = false;
static const bool WITH_NL = false;
@ -141,8 +144,9 @@ static const int INPUT_W = 128;
static const int OUTPUT_SIZE = 512;
static const int DEVICE_ID = 0;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const FastreidBackboneType BACKBONE = FastreidBackboneType::r34_distill;
static const FastreidHeadType HEAD = FastreidHeadType::EmbeddingHead;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const int LAST_STRIDE = 1;
static const bool WITH_IBNA = false;
static const bool WITH_NL = false;

19
projects/FastRT/demo/inference.cpp
View File

@ -3,8 +3,7 @@
#include "fastrt/utils.h"
#include "fastrt/baseline.h"
#include "fastrt/sbs_resnet.h"
#include "fastrt/embedding_head.h"
#include "fastrt/factory.h"
using namespace fastrt;
using namespace nvinfer1;
@ -18,8 +17,9 @@ static const int INPUT_W = 128;
static const int OUTPUT_SIZE = 2048;
static const int DEVICE_ID = 0;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const FastreidBackboneType BACKBONE = FastreidBackboneType::r50;
static const FastreidHeadType HEAD = FastreidHeadType::EmbeddingHead;
static const FastreidPoolingType HEAD_POOLING = FastreidPoolingType::gempoolP;
static const int LAST_STRIDE = 1;
static const bool WITH_IBNA = true;
static const bool WITH_NL = true;
@ -37,8 +37,9 @@ int main(int argc, char** argv) {
DEVICE_ID};
FastreidConfig reidCfg {
HEAD_POOLING,
BACKBONE,
HEAD,
HEAD_POOLING,
LAST_STRIDE,
WITH_IBNA,
WITH_NL,
@ -50,13 +51,11 @@ int main(int argc, char** argv) {
Baseline baseline{modelCfg, reidCfg};
if (argc == 2 && std::string(argv[1]) == "-s") {
auto backbone = createBackbone<IActivationLayer>(reidCfg);
if(!backbone) {
std::cout << "CreateBackbone Failed." << std::endl;
return -1;
}
ModuleFactory moduleFactory;
std::cout << "[Serializling Engine]" << std::endl;
if(!baseline.serializeEngine<IActivationLayer, IScaleLayer>(ENGINE_PATH, backbone, embedding_head)) {
if (!baseline.serializeEngine(ENGINE_PATH,
{std::move(moduleFactory.createBackbone(reidCfg.backbone)),
std::move(moduleFactory.createHead(reidCfg.head))})) {
std::cout << "SerializeEngine Failed." << std::endl;
return -1;
}

1
projects/FastRT/fastrt/CMakeLists.txt
View File

@ -49,3 +49,4 @@ add_subdirectory(engine)
add_subdirectory(heads)
add_subdirectory(backbones)
add_subdirectory(meta_arch)
add_subdirectory(factory)

9
projects/FastRT/fastrt/backbones/sbs_resnet.cpp
View File

@ -3,12 +3,11 @@
#include "fastrt/utils.h"
#include "fastrt/layers.h"
#include "fastrt/sbs_resnet.h"
using namespace trtxapi;
namespace fastrt {
IActivationLayer* backbone_sbsR34_distill(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
ILayer* backbone_sbsR34_distill::topology(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
std::string ibn{""};
if(reidCfg.with_ibna) {
ibn = "a";
@ -64,7 +63,7 @@ namespace fastrt {
return relu2;
}
IActivationLayer* backbone_sbsR50_distill(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
ILayer* backbone_sbsR50_distill::topology(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
std::string ibn{""};
if(reidCfg.with_ibna) {
ibn = "a";
@ -140,7 +139,7 @@ namespace fastrt {
return relu2;
}
IActivationLayer* backbone_sbsR34(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
ILayer* backbone_sbsR34::topology(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
std::string ibn{""};
if(reidCfg.with_ibna) {
ibn = "a";
@ -193,7 +192,7 @@ namespace fastrt {
return x;
}
IActivationLayer* backbone_sbsR50(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
ILayer* backbone_sbsR50::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/backbones/resnet.py
* NL layers follow by: nl_layers_per_stage = {'50x': [0, 2, 3, 0],}[depth]

26
projects/FastRT/fastrt/common/utils.cpp
View File

@ -67,13 +67,6 @@ namespace trt {
namespace fastrt {
const std::string PoolingTypetoString(FastreidPoolingType value) {
#define X(a, b) b,
static std::vector<std::string> table{ FASTPOOLING_TABLE };
#undef X
return table[value];
}
const std::string BackboneTypetoString(FastreidBackboneType value) {
#define X(a, b) b,
static std::vector<std::string> table{ FASTBACKBONE_TABLE };
@ -81,9 +74,24 @@ namespace fastrt {
return table[value];
}
const std::string HeadTypetoString(FastreidHeadType value) {
#define X(a, b) b,
static std::vector<std::string> table{ FASTHEAD_TABLE };
#undef X
return table[value];
}
const std::string PoolingTypetoString(FastreidPoolingType value) {
#define X(a, b) b,
static std::vector<std::string> table{ FASTPOOLING_TABLE };
#undef X
return table[value];
}
std::ostream& operator<<(std::ostream& os, const FastreidConfig& fastreidCfg) {
os << "\tpooling: " << PoolingTypetoString(fastreidCfg.pooling) << "\n\t"
<< "backbone: " << BackboneTypetoString(fastreidCfg.backbone) << "\n\t"
os << "\tbackbone: " << BackboneTypetoString(fastreidCfg.backbone) << "\n\t"
<< "head: " << HeadTypetoString(fastreidCfg.head) << "\n\t"
<< "pooling: " << PoolingTypetoString(fastreidCfg.pooling) << "\n\t"
<< "last_stride: " << fastreidCfg.last_stride << "\n\t"
<< "with_ibna: " << fastreidCfg.with_ibna << "\n\t"
<< "with_nl: " << fastreidCfg.with_nl << "\n\t"

4
projects/FastRT/fastrt/factory/CMakeLists.txt 100644
View File

@ -0,0 +1,4 @@
target_sources(${PROJECT_NAME}
PRIVATE
${CMAKE_CURRENT_SOURCE_DIR}/factory.cpp
)

49
projects/FastRT/fastrt/factory/factory.cpp 100644
View File

@ -0,0 +1,49 @@
#include <iostream>
#include "fastrt/utils.h"
#include "fastrt/sbs_resnet.h"
#include "fastrt/factory.h"
#include "fastrt/embedding_head.h"
namespace fastrt {
std::unique_ptr<Module> ModuleFactory::createBackbone(const FastreidBackboneType& backbonetype) {
switch(backbonetype) {
case FastreidBackboneType::r50:
/* cfg.MODEL.META_ARCHITECTURE: Baseline */
/* cfg.MODEL.BACKBONE.DEPTH: 50x */
std::cout << "[createBackboneModule]: backbone_sbsR50" << std::endl;
return make_unique<backbone_sbsR50>();
case FastreidBackboneType::r50_distill:
/* cfg.MODEL.META_ARCHITECTURE: Distiller */
/* cfg.MODEL.BACKBONE.DEPTH: 50x */
std::cout << "[createBackboneModule]: backbone_sbsR50_distill" << std::endl;
return make_unique<backbone_sbsR50_distill>();
case FastreidBackboneType::r34:
/* cfg.MODEL.META_ARCHITECTURE: Baseline */
/* cfg.MODEL.BACKBONE.DEPTH: 34x */
std::cout << "[createBackboneModule]: backbone_sbsR34" << std::endl;
return make_unique<backbone_sbsR34>();
case FastreidBackboneType::r34_distill:
/* cfg.MODEL.META_ARCHITECTURE: Distiller */
/* cfg.MODEL.BACKBONE.DEPTH: 34x */
std::cout << "[createBackboneModule]: backbone_sbsR34_distill" << std::endl;
return make_unique<backbone_sbsR34_distill>();
default:
std::cerr << "[Backbone is not supported.]" << std::endl;
return nullptr;
}
}
std::unique_ptr<Module> ModuleFactory::createHead(const FastreidHeadType& headtype) {
switch(headtype) {
case FastreidHeadType::EmbeddingHead:
/* cfg.MODEL.HEADS.NAME: EmbeddingHead */
std::cout << "[createHeadModule]: EmbeddingHead" << std::endl;
return make_unique<embedding_head>();
default:
std::cerr << "[Head is not supported.]" << std::endl;
return nullptr;
}
}
}

2
projects/FastRT/fastrt/heads/embedding_head.cpp
View File

@ -5,7 +5,7 @@
namespace fastrt {
IScaleLayer* embedding_head(INetworkDefinition *network, std::map<std::string, Weights>& weightMap, ITensor& input, const FastreidConfig& reidCfg) {
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
*/

64
projects/FastRT/fastrt/meta_arch/model.cpp
View File

@ -18,6 +18,70 @@ namespace fastrt {
_engineCfg.stream_size = 0;
};
bool Model::serializeEngine(const std::string engine_file, const std::initializer_list<std::unique_ptr<Module>>& modules) {
/* Create builder */
auto builder = make_holder(createInferBuilder(gLogger));
/* Create model to populate the network, then set the outputs and create an engine */
auto engine = createEngine(builder.get(), modules);
TRTASSERT(engine.get());
/* Serialize the engine */
auto modelStream = make_holder(engine->serialize());
TRTASSERT(modelStream.get());
std::ofstream p(engine_file, std::ios::binary | std::ios::out);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return false;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
std::cout << "[Save serialized engine]: " << engine_file << std::endl;
return true;
}
TensorRTHolder<ICudaEngine> Model::createEngine(IBuilder* builder, const std::initializer_list<std::unique_ptr<Module>>& modules) {
auto network = make_holder(builder->createNetworkV2(0U));
auto config = make_holder(builder->createBuilderConfig());
auto data = network->addInput(_engineCfg.input_name.c_str(), _dt, Dims3{3, _engineCfg.input_h, _engineCfg.input_w});
TRTASSERT(data);
auto weightMap = loadWeights(_engineCfg.weights_path);
/* Preprocessing */
auto input = preprocessing_gpu(network.get(), weightMap, data);
if (!input) input = data;
/* Modeling */
ILayer* output{nullptr};
for(auto& sequential_module: modules) {
output = sequential_module->topology(network.get(), weightMap, *input, _reidcfg);
TRTASSERT(output);
input = output->getOutput(0);
}
/* Set output */
output->getOutput(0)->setName(_engineCfg.output_name.c_str());
network->markOutput(*output->getOutput(0));
/* Build engine */
builder->setMaxBatchSize(_engineCfg.max_batch_size);
config->setMaxWorkspaceSize(1 << 20);
#ifdef BUILD_FP16
std::cout << "[Build fp16]" << std::endl;
config->setFlag(BuilderFlag::kFP16);
#endif
auto engine = make_holder(builder->buildEngineWithConfig(*network, *config));
std::cout << "[TRT engine build out]" << std::endl;
for (auto& mem : weightMap) {
free((void*) (mem.second.values));
}
return engine;
}
bool Model::deserializeEngine(const std::string engine_file) {
std::ifstream file(engine_file, std::ios::binary | std::ios::in);
if (file.good()) {

14
projects/FastRT/include/fastrt/embedding_head.h
View File

@ -2,14 +2,20 @@
#include <map>
#include "NvInfer.h"
#include "fastrt/module.h"
#include "fastrt/struct.h"
using namespace nvinfer1;
namespace fastrt {
IScaleLayer* embedding_head(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg);
class embedding_head : public Module {
public:
embedding_head() = default;
~embedding_head() = default;
ILayer* topology(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg) override;
};
}

17
projects/FastRT/include/fastrt/factory.h 100644
View File

@ -0,0 +1,17 @@
#pragma once
#include "struct.h"
#include "module.h"
namespace fastrt {
class ModuleFactory {
public:
ModuleFactory() = default;
~ModuleFactory() = default;
std::unique_ptr<Module> createBackbone(const FastreidBackboneType& backbonetype);
std::unique_ptr<Module> createHead(const FastreidHeadType& headtype);
};
}

74
projects/FastRT/include/fastrt/model.h
View File

@ -1,5 +1,6 @@
#pragma once
#include "module.h"
#include "utils.h"
#include "holder.h"
#include "layers.h"
@ -22,33 +23,15 @@ namespace fastrt {
const std::string input_name="input",
const std::string output_name="output");
virtual ~Model() {};
virtual ~Model() = default;
template <typename B, typename H>
/*
* Serialize TRT Engine
* @engine_file: save serialized engine as engine_file
* @modules: sequential modules(variadic length). (e.g., backbone1 + backbone2 + head, backbone + head, backbone)
*/
bool serializeEngine(const std::string engine_file,
std::function<B*(INetworkDefinition*, std::map<std::string, Weights>&, ITensor&, const FastreidConfig&)> backbone,
std::function<H*(INetworkDefinition*, std::map<std::string, Weights>&, ITensor&, const FastreidConfig&)> head) {
/* Create builder */
auto builder = make_holder(createInferBuilder(gLogger));
/* Create model to populate the network, then set the outputs and create an engine */
auto engine = createEngine<B, H>(builder.get(), backbone, head);
TRTASSERT(engine.get());
/* Serialize the engine */
auto modelStream = make_holder(engine->serialize());
TRTASSERT(modelStream.get());
std::ofstream p(engine_file, std::ios::binary | std::ios::out);
if (!p) {
std::cerr << "could not open plan output file" << std::endl;
return false;
}
p.write(reinterpret_cast<const char*>(modelStream->data()), modelStream->size());
std::cout << "[Save serialized engine]: " << engine_file << std::endl;
return true;
}
const std::initializer_list<std::unique_ptr<Module>>& modules);
bool deserializeEngine(const std::string engine_file);
@ -73,47 +56,8 @@ namespace fastrt {
int getDeviceID();
private:
template <typename B, typename H>
TensorRTHolder<ICudaEngine> createEngine(IBuilder* builder,
std::function<B*(INetworkDefinition*, std::map<std::string, Weights>&, ITensor&, const FastreidConfig&)> backbone,
std::function<H*(INetworkDefinition*, std::map<std::string, Weights>&, ITensor&, const FastreidConfig&)> head) {
auto network = make_holder(builder->createNetworkV2(0U));
auto config = make_holder(builder->createBuilderConfig());
auto data = network->addInput(_engineCfg.input_name.c_str(), _dt, Dims3{3, _engineCfg.input_h, _engineCfg.input_w});
TRTASSERT(data);
auto weightMap = loadWeights(_engineCfg.weights_path);
/* Preprocessing */
auto pre_input = preprocessing_gpu(network.get(), weightMap, data);
if (!pre_input) pre_input = data;
/* Modeling */
auto feat_map = backbone(network.get(), weightMap, *pre_input, _reidcfg);
TRTASSERT(feat_map);
auto embedding = head(network.get(), weightMap, *feat_map->getOutput(0), _reidcfg);
TRTASSERT(embedding);
/* Set output */
embedding->getOutput(0)->setName(_engineCfg.output_name.c_str());
network->markOutput(*embedding->getOutput(0));
/* Build engine */
builder->setMaxBatchSize(_engineCfg.max_batch_size);
config->setMaxWorkspaceSize(1 << 20);
#ifdef BUILD_FP16
std::cout << "[Build fp16]" << std::endl;
config->setFlag(BuilderFlag::kFP16);
#endif
auto engine = make_holder(builder->buildEngineWithConfig(*network, *config));
std::cout << "[TRT engine build out]" << std::endl;
for (auto& mem : weightMap) {
free((void*) (mem.second.values));
}
return engine;
}
const std::initializer_list<std::unique_ptr<Module>>& modules);
virtual void preprocessing_cpu(const cv::Mat& img, float* const data, const std::size_t stride) = 0;
virtual ITensor* preprocessing_gpu(INetworkDefinition* network,

21
projects/FastRT/include/fastrt/module.h 100644
View File

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

96
projects/FastRT/include/fastrt/sbs_resnet.h
View File

@ -1,71 +1,51 @@
#pragma once
#include <map>
#include <functional>
#include "struct.h"
#include "module.h"
#include "NvInfer.h"
using namespace nvinfer1;
namespace fastrt {
IActivationLayer* backbone_sbsR34_distill(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg);
class backbone_sbsR34_distill : public Module {
public:
backbone_sbsR34_distill() = default;
~backbone_sbsR34_distill() = default;
ILayer* topology(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg) override;
};
IActivationLayer* backbone_sbsR50_distill(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg);
class backbone_sbsR50_distill : public Module {
public:
backbone_sbsR50_distill() = default;
~backbone_sbsR50_distill() = default;
ILayer* topology(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg) override;
};
IActivationLayer* backbone_sbsR34(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg);
IActivationLayer* backbone_sbsR50(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg);
}
namespace fastrt {
template <typename T>
using backboneFcn = std::function<T*(INetworkDefinition*, std::map<std::string, Weights>&, ITensor&, const FastreidConfig&)>;
template <typename T>
backboneFcn<T> createBackbone(const FastreidConfig &reidcfg) {
switch(reidcfg.backbone) {
case FastreidBackboneType::r50:
/* cfg.MODEL.META_ARCHITECTURE: Baseline */
/* cfg.MODEL.BACKBONE.DEPTH: 50x */
std::cout << "[CreateBackbone]: backbone_sbsR50" << std::endl;
return backbone_sbsR50;
break;
case FastreidBackboneType::r50_distill:
/* cfg.MODEL.META_ARCHITECTURE: Distiller */
/* cfg.MODEL.BACKBONE.DEPTH: 50x */
std::cout << "[CreateBackbone]: backbone_sbsR50_distill" << std::endl;
return backbone_sbsR50_distill;
break;
case FastreidBackboneType::r34:
/* cfg.MODEL.META_ARCHITECTURE: Baseline */
/* cfg.MODEL.BACKBONE.DEPTH: 34x */
std::cout << "[CreateBackbone]: backbone_sbsR34" << std::endl;
return backbone_sbsR34;
break;
case FastreidBackboneType::r34_distill:
/* cfg.MODEL.META_ARCHITECTURE: Distiller */
/* cfg.MODEL.BACKBONE.DEPTH: 34x */
std::cout << "[CreateBackbone]: backbone_sbsR34_distill" << std::endl;
return backbone_sbsR34_distill;
break;
default:
std::cout << "[Backbone is not supported.]" << std::endl;
}
return nullptr;
}
class backbone_sbsR34 : public Module {
public:
backbone_sbsR34() = default;
~backbone_sbsR34() = default;
ILayer* topology(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg) override;
};
class backbone_sbsR50 : public Module {
public:
backbone_sbsR50() = default;
~backbone_sbsR50() = default;
ILayer* topology(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input,
const FastreidConfig& reidCfg) override;
};
}

30
projects/FastRT/include/fastrt/struct.h
View File

@ -24,16 +24,6 @@ namespace trt {
namespace fastrt {
#define FASTPOOLING_TABLE \
X(maxpool, "maxpool") \
X(avgpool, "avgpool") \
X(gempool, "gempool") \
X(gempoolP, "gempoolP")
#define X(a, b) a,
enum FastreidPoolingType { FASTPOOLING_TABLE };
#undef X
#define FASTBACKBONE_TABLE \
X(r50, "r50") \
X(r50_distill, "r50_distill") \
@ -44,9 +34,27 @@ namespace fastrt {
enum FastreidBackboneType { FASTBACKBONE_TABLE };
#undef X
#define FASTHEAD_TABLE \
X(EmbeddingHead, "EmbeddingHead")
#define X(a, b) a,
enum FastreidHeadType { FASTHEAD_TABLE };
#undef X
#define FASTPOOLING_TABLE \
X(maxpool, "maxpool") \
X(avgpool, "avgpool") \
X(gempool, "gempool") \
X(gempoolP, "gempoolP")
#define X(a, b) a,
enum FastreidPoolingType { FASTPOOLING_TABLE };
#undef X
struct FastreidConfig {
FastreidPoolingType pooling; /* cfg.MODEL.HEADS.POOL_LAYER */
FastreidBackboneType backbone; /* cfg.MODEL.BACKBONE.DEPTH and cfg.MODEL.META_ARCHITECTURE */
FastreidHeadType head; /* cfg.MODEL.HEADS.NAME */
FastreidPoolingType pooling; /* cfg.MODEL.HEADS.POOL_LAYER */
int last_stride; /* cfg.MODEL.BACKBONE.LAST_STRIDE */
bool with_ibna; /* cfg.MODEL.BACKBONE.WITH_IBN */
bool with_nl; /* cfg.MODEL.BACKBONE.WITH_NL */

15
projects/FastRT/include/fastrt/utils.h
View File

@ -6,8 +6,8 @@
#include <vector>
#include <fstream>
#include <iostream>
#include <cassert>
#include "assert.h"
#include "NvInfer.h"
#include "cuda_runtime_api.h"
#include "fastrt/struct.h"
@ -23,18 +23,7 @@
} \
} while (0)
#define TRTASSERT(CONDITION) \
do \
{ \
auto cond = (CONDITION); \
if (!cond) \
{ \
std::cerr << "Condition failure.\n"; \
abort(); \
} \
} while (0)
#define TRTASSERT assert
using Time = std::chrono::high_resolution_clock;
using TimePoint = std::chrono::time_point<std::chrono::high_resolution_clock>;