mirror of https://github.com/JDAI-CV/fast-reid.git
64 lines
1.8 KiB
C++
64 lines
1.8 KiB
C++
#pragma once
|
|
|
|
#include <memory>
|
|
|
|
namespace trt {
|
|
|
|
struct ModelConfig {
|
|
std::string weights_path;
|
|
int max_batch_size;
|
|
int input_h; /* cfg.INPUT.SIZE_TRAIN[0] */
|
|
int input_w; /* cfg.INPUT.SIZE_TRAIN[1] */
|
|
int output_size; /* final embedding dims. Could be cfg.MODEL.BACKBONE.FEAT_DIM or cfg.MODEL.HEADS.EMBEDDING_DIM(if you modified. default=0) */
|
|
int device_id; /* cuda device id(0, 1, 2, ...) */
|
|
};
|
|
|
|
struct EngineConfig : ModelConfig {
|
|
std::string input_name;
|
|
std::string output_name;
|
|
std::shared_ptr<char> trtModelStream;
|
|
int stream_size;
|
|
};
|
|
|
|
}
|
|
|
|
namespace fastrt {
|
|
|
|
#define FASTBACKBONE_TABLE \
|
|
X(r50, "r50") \
|
|
X(r50_distill, "r50_distill") \
|
|
X(r34, "r34") \
|
|
X(r34_distill, "r34_distill")
|
|
|
|
#define X(a, b) a,
|
|
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 {
|
|
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 */
|
|
int embedding_dim; /* cfg.MODEL.HEADS.EMBEDDING_DIM (Default = 0) */
|
|
};
|
|
|
|
} |