wehub-resource-sync
55 lines
2.2 KiB
C++
55 lines
2.2 KiB
C++
//
|
|
// TransformBatchNormal.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/05.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../PostTreatUtils.hpp"
|
|
|
|
class TransformBatchNormal : public PostConverter {
|
|
public:
|
|
virtual bool onExecute(std::unique_ptr<MNN::NetT>& net) const override {
|
|
for (auto iter = net->oplists.begin(); iter != net->oplists.end();) {
|
|
auto& op = *iter;
|
|
const MNN::OpType opType = op->type;
|
|
|
|
if (MNN::OpType_BatchNorm != opType) {
|
|
iter++;
|
|
continue;
|
|
}
|
|
const int inputSize = op->inputIndexes.size();
|
|
DCHECK(inputSize == 1 || inputSize == 3) << "MNN BatchnNorm input size error!";
|
|
// instance norm have three input tensors(input_tensor, mean, variance)
|
|
if (inputSize == 3) {
|
|
iter++;
|
|
continue;
|
|
}
|
|
// DLOG(INFO) << "change BatchNorm to Scale: " << op->name;
|
|
auto batchnormParam = op->main.AsBatchNorm();
|
|
auto scaleParam = new MNN::ScaleT;
|
|
scaleParam->channels = batchnormParam->channels;
|
|
scaleParam->scaleData.resize(batchnormParam->channels);
|
|
scaleParam->biasData.resize(batchnormParam->channels);
|
|
const float* slopePtr = batchnormParam->slopeData.data();
|
|
const float* meanDataPtr = batchnormParam->meanData.data();
|
|
const float* varDataPtr = batchnormParam->varData.data();
|
|
const float* biasDataPtr = batchnormParam->biasData.data();
|
|
const float eps = batchnormParam->epsilon;
|
|
|
|
for (int i = 0; i < batchnormParam->channels; i++) {
|
|
float sqrt_var = sqrt(varDataPtr[i] + eps);
|
|
scaleParam->biasData[i] = biasDataPtr[i] - slopePtr[i] * meanDataPtr[i] / sqrt_var;
|
|
scaleParam->scaleData[i] = slopePtr[i] / sqrt_var;
|
|
}
|
|
|
|
op->type = MNN::OpType_Scale;
|
|
op->main.type = MNN::OpParameter_Scale;
|
|
op->main.value = scaleParam;
|
|
}
|
|
return true;
|
|
}
|
|
};
|
|
static PostConverterRegister<TransformBatchNormal> __l("TransformBatchNormal");
|