36 lines
1.1 KiB
C++
36 lines
1.1 KiB
C++
/*
|
|
* Created by Marcos Luciano
|
|
* https://www.github.com/marcoslucianops
|
|
*/
|
|
|
|
#include "maxpool_layer.h"
|
|
|
|
nvinfer1::ILayer* maxpoolLayer(
|
|
int layerIdx,
|
|
std::map<std::string, std::string>& block,
|
|
nvinfer1::ITensor* input,
|
|
nvinfer1::INetworkDefinition* network)
|
|
{
|
|
assert(block.at("type") == "maxpool");
|
|
assert(block.find("size") != block.end());
|
|
assert(block.find("stride") != block.end());
|
|
|
|
int size = std::stoi(block.at("size"));
|
|
int stride = std::stoi(block.at("stride"));
|
|
|
|
nvinfer1::IPoolingLayer* pool
|
|
= network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX, nvinfer1::Dims{2, {size, size}});
|
|
assert(pool);
|
|
std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
|
|
pool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
|
|
pool->setPaddingNd(nvinfer1::Dims{2, {(size - 1) / 2, (size - 1) / 2}});
|
|
if (size == 2 && stride == 1)
|
|
{
|
|
pool->setPrePadding(nvinfer1::Dims{2, {0, 0}});
|
|
pool->setPostPadding(nvinfer1::Dims{2, {1, 1}});
|
|
}
|
|
pool->setName(maxpoolLayerName.c_str());
|
|
|
|
return pool;
|
|
}
|