darrenhsieh 0617e3eeb7 refactor: embedding_head
* add abstract poolinglayer

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

46 lines
1.2 KiB
C++

#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<std::string, Weights>& weightMap,
ITensor& input) override;
};
class AvgPool : public IPoolingLayerRT {
public:
AvgPool() = default;
~AvgPool() = default;
ILayer* addPooling(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input) override;
};
class GemPool : public IPoolingLayerRT {
public:
GemPool() = default;
~GemPool() = default;
ILayer* addPooling(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input) override;
};
class GemPoolP : public IPoolingLayerRT {
public:
GemPoolP() = default;
~GemPoolP() = default;
ILayer* addPooling(INetworkDefinition *network,
std::map<std::string, Weights>& weightMap,
ITensor& input) override;
};
}