55 lines
1.5 KiB
C++
55 lines
1.5 KiB
C++
/*
|
|
* Created by Marcos Luciano
|
|
* https://www.github.com/marcoslucianops
|
|
*/
|
|
|
|
#include "reduce_layer.h"
|
|
|
|
nvinfer1::ITensor*
|
|
reduceLayer(int layerIdx, std::map<std::string, std::string>& block, nvinfer1::ITensor* input,
|
|
nvinfer1::INetworkDefinition* network)
|
|
{
|
|
nvinfer1::ITensor* output;
|
|
|
|
assert(block.at("type") == "reduce");
|
|
assert(block.find("mode") != block.end());
|
|
assert(block.find("axes") != block.end());
|
|
|
|
std::string mode = block.at("mode");
|
|
|
|
nvinfer1::ReduceOperation operation;
|
|
if (mode == "mean")
|
|
operation = nvinfer1::ReduceOperation::kAVG;
|
|
|
|
std::string strAxes = block.at("axes");
|
|
std::vector<int32_t> axes;
|
|
size_t lastPos = 0, pos = 0;
|
|
while ((pos = strAxes.find(',', lastPos)) != std::string::npos) {
|
|
int vL = std::stoi(trim(strAxes.substr(lastPos, pos - lastPos)));
|
|
axes.push_back(vL);
|
|
lastPos = pos + 1;
|
|
}
|
|
if (lastPos < strAxes.length()) {
|
|
std::string lastV = trim(strAxes.substr(lastPos));
|
|
if (!lastV.empty())
|
|
axes.push_back(std::stoi(lastV));
|
|
}
|
|
assert(!axes.empty());
|
|
|
|
uint32_t axisMask = 0;
|
|
for (int axis : axes)
|
|
axisMask |= 1 << axis;
|
|
|
|
bool keepDims = false;
|
|
if (block.find("keep") != block.end())
|
|
keepDims = std::stoi(block.at("keep")) == 1 ? true : false;
|
|
|
|
nvinfer1::IReduceLayer* reduce = network->addReduce(*input, operation, axisMask, keepDims);
|
|
assert(reduce != nullptr);
|
|
std::string reduceLayerName = "reduce_" + std::to_string(layerIdx);
|
|
reduce->setName(reduceLayerName.c_str());
|
|
output = reduce->getOutput(0);
|
|
|
|
return output;
|
|
}
|