55 lines
1.8 KiB
C++
55 lines
1.8 KiB
C++
/*
|
|
* Created by Marcos Luciano
|
|
* https://www.github.com/marcoslucianops
|
|
*/
|
|
|
|
#include "pooling_layer.h"
|
|
|
|
#include <cassert>
|
|
#include <iostream>
|
|
|
|
nvinfer1::ITensor*
|
|
poolingLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::ITensor* input,
|
|
nvinfer1::INetworkDefinition* network)
|
|
{
|
|
nvinfer1::ITensor* output;
|
|
|
|
assert(block.at("type") == "maxpool" || block.at("type") == "avgpool");
|
|
|
|
if (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* maxpool = network->addPoolingNd(*input, nvinfer1::PoolingType::kMAX,
|
|
nvinfer1::Dims{2, {size, size}});
|
|
assert(maxpool != nullptr);
|
|
std::string maxpoolLayerName = "maxpool_" + std::to_string(layerIdx);
|
|
maxpool->setName(maxpoolLayerName.c_str());
|
|
maxpool->setStrideNd(nvinfer1::Dims{2, {stride, stride}});
|
|
maxpool->setPaddingNd(nvinfer1::Dims{2, {(size - 1) / 2, (size - 1) / 2}});
|
|
if (size == 2 && stride == 1) {
|
|
maxpool->setPrePadding(nvinfer1::Dims{2, {0, 0}});
|
|
maxpool->setPostPadding(nvinfer1::Dims{2, {1, 1}});
|
|
}
|
|
output = maxpool->getOutput(0);
|
|
}
|
|
else if (block.at("type") == "avgpool") {
|
|
nvinfer1::Dims inputDims = input->getDimensions();
|
|
nvinfer1::IPoolingLayer* avgpool = network->addPoolingNd(*input, nvinfer1::PoolingType::kAVERAGE,
|
|
nvinfer1::Dims{2, {inputDims.d[1], inputDims.d[2]}});
|
|
assert(avgpool != nullptr);
|
|
std::string avgpoolLayerName = "avgpool_" + std::to_string(layerIdx);
|
|
avgpool->setName(avgpoolLayerName.c_str());
|
|
output = avgpool->getOutput(0);
|
|
}
|
|
else {
|
|
std::cerr << "Pooling not supported: " << block.at("type") << std::endl;
|
|
assert(0);
|
|
}
|
|
|
|
return output;
|
|
}
|