mirror of
https://github.com/JDAI-CV/fast-reid.git
synced 2025-06-03 14:50:47 +08:00
27 lines
675 B
C++
27 lines
675 B
C++
#pragma once
|
|
|
|
#include <map>
|
|
#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:
|
|
FastreidConfig& _modelCfg;
|
|
std::unique_ptr<LayerFactory> _layerFactory;
|
|
|
|
public:
|
|
embedding_head(FastreidConfig& modelCfg);
|
|
embedding_head(FastreidConfig& modelCfg, std::unique_ptr<LayerFactory> layerFactory);
|
|
~embedding_head() = default;
|
|
|
|
ILayer* topology(INetworkDefinition *network,
|
|
std::map<std::string, Weights>& weightMap,
|
|
ITensor& input) override;
|
|
};
|
|
|
|
} |