115 lines
3.2 KiB
C++
115 lines
3.2 KiB
C++
/*
|
|
* Created by Marcos Luciano
|
|
* https://www.github.com/marcoslucianops
|
|
*/
|
|
|
|
#include <math.h>
|
|
#include "batchnorm_layer.h"
|
|
|
|
nvinfer1::ILayer* batchnormLayer(
|
|
int layerIdx,
|
|
std::map<std::string, std::string>& block,
|
|
std::vector<float>& weights,
|
|
std::vector<nvinfer1::Weights>& trtWeights,
|
|
int& weightPtr,
|
|
std::string weightsType,
|
|
float eps,
|
|
nvinfer1::ITensor* input,
|
|
nvinfer1::INetworkDefinition* network)
|
|
{
|
|
assert(block.at("type") == "batchnorm");
|
|
assert(block.find("filters") != block.end());
|
|
|
|
int filters = std::stoi(block.at("filters"));
|
|
std::string activation = block.at("activation");
|
|
|
|
std::vector<float> bnBiases;
|
|
std::vector<float> bnWeights;
|
|
std::vector<float> bnRunningMean;
|
|
std::vector<float> bnRunningVar;
|
|
|
|
if (weightsType == "weights") {
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnBiases.push_back(weights[weightPtr]);
|
|
weightPtr++;
|
|
}
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnWeights.push_back(weights[weightPtr]);
|
|
weightPtr++;
|
|
}
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnRunningMean.push_back(weights[weightPtr]);
|
|
weightPtr++;
|
|
}
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnRunningVar.push_back(sqrt(weights[weightPtr] + 1.0e-5));
|
|
weightPtr++;
|
|
}
|
|
}
|
|
else {
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnWeights.push_back(weights[weightPtr]);
|
|
weightPtr++;
|
|
}
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnBiases.push_back(weights[weightPtr]);
|
|
weightPtr++;
|
|
}
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnRunningMean.push_back(weights[weightPtr]);
|
|
weightPtr++;
|
|
}
|
|
for (int i = 0; i < filters; ++i)
|
|
{
|
|
bnRunningVar.push_back(sqrt(weights[weightPtr] + eps));
|
|
weightPtr++;
|
|
}
|
|
}
|
|
|
|
int size = filters;
|
|
nvinfer1::Weights shift{nvinfer1::DataType::kFLOAT, nullptr, size};
|
|
nvinfer1::Weights scale{nvinfer1::DataType::kFLOAT, nullptr, size};
|
|
nvinfer1::Weights power{nvinfer1::DataType::kFLOAT, nullptr, size};
|
|
float* shiftWt = new float[size];
|
|
for (int i = 0; i < size; ++i)
|
|
{
|
|
shiftWt[i]
|
|
= bnBiases.at(i) - ((bnRunningMean.at(i) * bnWeights.at(i)) / bnRunningVar.at(i));
|
|
}
|
|
shift.values = shiftWt;
|
|
float* scaleWt = new float[size];
|
|
for (int i = 0; i < size; ++i)
|
|
{
|
|
scaleWt[i] = bnWeights.at(i) / bnRunningVar[i];
|
|
}
|
|
scale.values = scaleWt;
|
|
float* powerWt = new float[size];
|
|
for (int i = 0; i < size; ++i)
|
|
{
|
|
powerWt[i] = 1.0;
|
|
}
|
|
power.values = powerWt;
|
|
trtWeights.push_back(shift);
|
|
trtWeights.push_back(scale);
|
|
trtWeights.push_back(power);
|
|
|
|
nvinfer1::IScaleLayer* bn = network->addScale(
|
|
*input, nvinfer1::ScaleMode::kCHANNEL, shift, scale, power);
|
|
assert(bn != nullptr);
|
|
std::string bnLayerName = "batch_norm_" + std::to_string(layerIdx);
|
|
bn->setName(bnLayerName.c_str());
|
|
nvinfer1::ILayer* output = bn;
|
|
|
|
output = activationLayer(layerIdx, activation, output, output->getOutput(0), network);
|
|
assert(output != nullptr);
|
|
|
|
return output;
|
|
}
|