fast-reid/projects/FastRT/include/fastrt/IPoolingLayerRT.h
darrenhsieh 0617e3eeb7 refactor: embedding_head
* add abstract poolinglayer

* add poolinglayer factory
2021-02-13 21:28:29 +08:00

20 lines
405 B
C++

#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;
};
}