wehub-resource-sync
377 lines
15 KiB
C++
377 lines
15 KiB
C++
//
|
|
// OnnxBatchNormMerge.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/10/16.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <math.h>
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
namespace MNN {
|
|
namespace Express {
|
|
static VARP _ReshapeF(VARP x, VARP shape, MNN::MNN_DATA_FORMAT format) {
|
|
MNN_ASSERT(nullptr != x);
|
|
std::unique_ptr<OpT> reshape(new OpT);
|
|
reshape->type = OpType_Reshape;
|
|
reshape->main.type = OpParameter_Reshape;
|
|
reshape->main.value = new ReshapeT;
|
|
reshape->main.AsReshape()->dimType = format;
|
|
return (Variable::create(Expr::create(reshape.get(), {x, shape})));
|
|
}
|
|
|
|
class OnnxBatchNormTransform : public OnnxExtraManager::Transform {
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
|
|
MNN_THROW_CHECK(inputs.size() == 5, "BatchNorm should have 5 inputs");
|
|
|
|
int channels = 1;
|
|
float epsilon = 1e-10;
|
|
|
|
auto bnOp = expr->get();
|
|
auto extraParam = bnOp->main_as_Extra();
|
|
int size = 0;
|
|
if (nullptr != extraParam->attr()) {
|
|
size = extraParam->attr()->size();
|
|
for (int i = 0; i < size; ++i) {
|
|
auto attr = extraParam->attr()->GetAs<Attribute>(i);
|
|
const auto& key = attr->key()->str();
|
|
if (key == "epsilon") {
|
|
epsilon = attr->f();
|
|
}
|
|
}
|
|
}
|
|
|
|
auto gamma = inputs[1];
|
|
auto beta = inputs[2];
|
|
auto mean = inputs[3];
|
|
auto variance = inputs[4];
|
|
|
|
MNN_THROW_CHECK(gamma->getInfo() != nullptr, "BatchNorm second input should be Constant!");
|
|
MNN_THROW_CHECK(beta->getInfo() != nullptr, "BatchNorm second input should be Constant!");
|
|
MNN_THROW_CHECK(mean->getInfo() != nullptr, "BatchNorm second input should be Constant!");
|
|
MNN_THROW_CHECK(variance->getInfo() != nullptr, "BatchNorm second input should be Constant!");
|
|
auto gammaSize = gamma->getInfo()->size;
|
|
auto betaSize = beta->getInfo()->size;
|
|
auto meanSize = mean->getInfo()->size;
|
|
auto varianceSize = variance->getInfo()->size;
|
|
|
|
// find the max value(incase broadcast mode)
|
|
channels = gammaSize > betaSize ? gammaSize : betaSize;
|
|
channels = channels > meanSize ? channels : meanSize;
|
|
channels = channels > varianceSize ? channels : varianceSize;
|
|
|
|
std::unique_ptr<MNN::BatchNormT> batchnorm(new MNN::BatchNormT);
|
|
batchnorm->slopeData.resize(channels);
|
|
batchnorm->biasData.resize(channels);
|
|
batchnorm->meanData.resize(channels);
|
|
batchnorm->varData.resize(channels);
|
|
batchnorm->channels = channels;
|
|
|
|
// TODO check data length, then support broadcast mode
|
|
auto gammaDataPtr = gamma->readMap<float>();
|
|
MNN_THROW_CHECK(gammaDataPtr != nullptr, "BatchNorm's gamma not valid!");
|
|
memcpy(batchnorm->slopeData.data(), gammaDataPtr, gamma->getInfo()->size * sizeof(float));
|
|
|
|
auto betaDataPtr = beta->readMap<float>();
|
|
MNN_THROW_CHECK(betaDataPtr != nullptr, "BatchNorm's beta not valid!");
|
|
memcpy(batchnorm->biasData.data(), betaDataPtr, beta->getInfo()->size * sizeof(float));
|
|
auto meanDataPtr = mean->readMap<float>();
|
|
MNN_THROW_CHECK(meanDataPtr != nullptr, "BatchNorm's mean not valid!");
|
|
memcpy(batchnorm->meanData.data(), meanDataPtr, mean->getInfo()->size * sizeof(float));
|
|
auto varPtr = variance->readMap<float>();
|
|
MNN_THROW_CHECK(varPtr != nullptr, "BatchNorm's var not valid!");
|
|
for (int i = 0; i < channels; ++i) {
|
|
batchnorm->varData[i] = varPtr[i];
|
|
}
|
|
|
|
std::unique_ptr<OpT> mnnBnOp(new OpT);
|
|
mnnBnOp->name = expr->name();
|
|
mnnBnOp->type = OpType_BatchNorm;
|
|
mnnBnOp->main.type = OpParameter_BatchNorm;
|
|
{
|
|
auto bnParam = new MNN::BatchNormT;
|
|
mnnBnOp->main.value = bnParam;
|
|
bnParam->channels = batchnorm->channels;
|
|
bnParam->slopeData.resize(batchnorm->channels);
|
|
bnParam->biasData.resize(batchnorm->channels);
|
|
bnParam->meanData.resize(batchnorm->channels);
|
|
bnParam->varData.resize(batchnorm->channels);
|
|
const float* slopeDataPtr = batchnorm->slopeData.data();
|
|
const float* biasDataPtr = batchnorm->biasData.data();
|
|
const float* meanDataPtr = batchnorm->meanData.data();
|
|
const float* varDataPtr = batchnorm->varData.data();
|
|
|
|
for (int i = 0; i < batchnorm->channels; i++) {
|
|
bnParam->slopeData[i] = slopeDataPtr[i];
|
|
bnParam->biasData[i] = biasDataPtr[i];
|
|
bnParam->meanData[i] = meanDataPtr[i];
|
|
bnParam->varData[i] = varDataPtr[i];
|
|
}
|
|
bnParam->epsilon = epsilon;
|
|
}
|
|
// create merged op
|
|
auto newExpr = Expr::create(mnnBnOp.get(), {inputs[0]});
|
|
newExpr->setName(expr->name());
|
|
auto res = Variable::create(newExpr);
|
|
return res->expr().first;
|
|
}
|
|
};
|
|
|
|
static VARP _OnnxReshape(VARP x, VARP shape) {
|
|
std::unique_ptr<OpT> reshape(new OpT);
|
|
reshape->type = OpType_Reshape;
|
|
reshape->main.type = OpParameter_Reshape;
|
|
reshape->main.value = new ReshapeT;
|
|
reshape->main.AsReshape()->dimType = MNN_DATA_FORMAT_NCHW;
|
|
return (Variable::create(Expr::create(reshape.get(), {x, shape})));
|
|
}
|
|
|
|
class OnnxInstanceNormalTransform : public OnnxExtraManager::Transform {
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
|
|
MNN_THROW_CHECK(inputs.size() == 3, "InstanceNormal should have 3 inputs");
|
|
auto input = inputs[0];
|
|
|
|
int channels = 1;
|
|
float epsilon = 1e-10;
|
|
|
|
auto bnOp = expr->get();
|
|
auto extraParam = bnOp->main_as_Extra();
|
|
int size = 0;
|
|
if (nullptr != extraParam->attr()) {
|
|
size = extraParam->attr()->size();
|
|
for (int i = 0; i < size; ++i) {
|
|
auto attr = extraParam->attr()->GetAs<Attribute>(i);
|
|
const auto& key = attr->key()->str();
|
|
if (key == "epsilon") {
|
|
epsilon = attr->f();
|
|
}
|
|
}
|
|
}
|
|
bool needScale = true;
|
|
bool scaleConst = false;
|
|
do {
|
|
auto biasPtr = inputs[2]->readMap<float>();
|
|
auto scalePtr = inputs[1]->readMap<float>();
|
|
if (nullptr == biasPtr || nullptr == scalePtr) {
|
|
break;
|
|
}
|
|
scaleConst = true;
|
|
auto oneVar = _Scalar<float>(1.0f);
|
|
auto scaleOff = inputs[1] - oneVar;
|
|
auto scaleSum = _ReduceSum(scaleOff * scaleOff);
|
|
if (scaleSum->readMap<float>()[0] > 0.000001f) {
|
|
break;
|
|
}
|
|
auto biasSum = _ReduceSum(inputs[2] * inputs[2]);
|
|
if (biasSum->readMap<float>()[0] > 0.000001f) {
|
|
break;
|
|
}
|
|
needScale = false;
|
|
} while (false);
|
|
auto originShape = _Shape(inputs[0], NCHW);
|
|
auto inputDim3 = _Reshape(inputs[0], {0, 0, -1}, NCHW);
|
|
|
|
// Turn to layernorm
|
|
std::unique_ptr<MNN::OpT> layerNormOp(new MNN::OpT);
|
|
layerNormOp->type = OpType_LayerNorm;
|
|
layerNormOp->main.value = new LayerNormT;
|
|
layerNormOp->main.type = OpParameter_LayerNorm;
|
|
{
|
|
auto param = layerNormOp->main.AsLayerNorm();
|
|
param->axis = {-1}; // Layernorm only need axis's size as 1
|
|
param->epsilon = epsilon;
|
|
param->group = 1;
|
|
}
|
|
auto res = Variable::create(Expr::create(layerNormOp.get(), {inputDim3}));
|
|
res = _ReshapeF(res, originShape, MNN_DATA_FORMAT_NCHW);
|
|
if (needScale) {
|
|
if (scaleConst) {
|
|
auto biasPtr = inputs[2]->readMap<float>();
|
|
auto scalePtr = inputs[1]->readMap<float>();
|
|
int channels = inputs[1]->getInfo()->size;
|
|
std::vector<float> scales(channels);
|
|
std::vector<float> bias(channels);
|
|
::memcpy(bias.data(), biasPtr, channels * sizeof(float));
|
|
::memcpy(scales.data(), scalePtr, channels * sizeof(float));
|
|
res = _Scale(res, channels, std::move(scales), std::move(bias));
|
|
} else {
|
|
auto compatShape = _Concat({_Shape(inputs[1], true), _Fill(_Unsqueeze(_Size(_Shape(input, true)) - _Scalar<int>(2), {0}), _Scalar<int>(1))}, 0);
|
|
auto scale = _OnnxReshape(inputs[1], compatShape);
|
|
auto bias = _OnnxReshape(inputs[2], compatShape);
|
|
res = res * scale + bias;
|
|
}
|
|
}
|
|
res->setName(expr->name());
|
|
return res->expr().first;
|
|
}
|
|
};
|
|
|
|
class OnnxMeanVarianceNormTransform : public OnnxExtraManager::Transform {
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
std::vector<int> axes {0, 2, 3};
|
|
auto attrs = expr->get()->main_as_Extra()->attr();
|
|
if (attrs != nullptr) {
|
|
for (const auto& attr : *attrs) {
|
|
if (attr->key()->str() == "axes") {
|
|
axes.clear();
|
|
for (int i = 0; i < attr->list()->i()->size(); ++i) {
|
|
axes.push_back(attr->list()->i()->Get(i));
|
|
}
|
|
}
|
|
}
|
|
}
|
|
auto input = expr->inputs()[0];
|
|
auto mean = _ReduceMean(input, axes, true);
|
|
auto temp = input - mean;
|
|
auto var = _ReduceMean(temp * temp, axes, true);
|
|
auto res = temp * _Rsqrt(var);
|
|
res->setName(expr->name());
|
|
return res->expr().first;
|
|
}
|
|
};
|
|
|
|
class OnnxLpNormTransform : public OnnxExtraManager::Transform {
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto input = expr->inputs()[0];
|
|
int p = 2, axis = -1;
|
|
auto attrs = expr->get()->main_as_Extra()->attr();
|
|
if (attrs != nullptr) {
|
|
for (const auto& attr : *attrs) {
|
|
auto attrName = attr->key()->str();
|
|
if (attrName == "axis") {
|
|
axis = attr->i();
|
|
} else if (attrName == "p") {
|
|
p = attr->i();
|
|
}
|
|
}
|
|
}
|
|
if (p != 1 && p != 2) {
|
|
MNN_ERROR("Onnx's LpNormalization only support attr p is 1 or 2");
|
|
return nullptr;
|
|
}
|
|
VARP res;
|
|
if (p == 1) {
|
|
res = input / _ReduceSumMutable(_Abs(input), _Scalar<int>(axis), true);
|
|
} else {
|
|
res = input * _Rsqrt(_ReduceSumMutable(input * input, _Scalar<int>(axis), true));
|
|
}
|
|
res->setName(expr->name());
|
|
return res->expr().first;
|
|
}
|
|
};
|
|
|
|
class OnnxLayerNormTransform : public OnnxExtraManager::Transform {
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto input = expr->inputs()[0];
|
|
int axis = -1;
|
|
float eps = 1e-05;
|
|
auto attrs = expr->get()->main_as_Extra()->attr();
|
|
if (attrs != nullptr) {
|
|
for (const auto& attr : *attrs) {
|
|
auto attrName = attr->key()->str();
|
|
if (attrName == "axis") {
|
|
axis = attr->i();
|
|
}
|
|
if (attrName == "epsilon") {
|
|
eps = attr->f();
|
|
}
|
|
}
|
|
}
|
|
if (expr->outputSize() > 1 || axis > 0) {
|
|
// If axis > 0, we can't determine how many axis should be norm
|
|
auto axisVar = _Scalar<int>(axis);
|
|
// Add negative protect, may decrease performance
|
|
auto rankVar = _Rank(inputs[0]);
|
|
axisVar = _Mod(axisVar + rankVar, rankVar);
|
|
auto reduceAxis = _Range(axisVar, rankVar, _Scalar<int>(1));
|
|
auto mean = _ReduceMeanMutable(input, reduceAxis, true);
|
|
auto sub = input - mean;
|
|
auto normal = _Rsqrt(_ReduceMeanMutable(_Square(sub), reduceAxis, true) + _Scalar<float>(eps));
|
|
auto y = sub * normal * inputs[1];
|
|
if (inputs.size() > 2) {
|
|
y = y + inputs[2];
|
|
}
|
|
std::vector<VARP> identityOutputs = {y};
|
|
if (expr->outputSize() > 1) {
|
|
identityOutputs.emplace_back(mean);
|
|
}
|
|
if (expr->outputSize() > 2) {
|
|
identityOutputs.emplace_back(normal);
|
|
}
|
|
std::unique_ptr<OpT> copyOp(new OpT);
|
|
copyOp->type = OpType_Identity;
|
|
auto resultExpr = Expr::create(copyOp.get(), identityOutputs, identityOutputs.size());
|
|
resultExpr->setName(expr->name());
|
|
for (int i=0; i<expr->outputSize(); ++i) {
|
|
auto var = MNN::Express::Variable::create(resultExpr, i);
|
|
var->setName(expr->outputName(i));
|
|
}
|
|
return resultExpr;
|
|
}
|
|
std::shared_ptr<MNN::OpT> layernorm(new MNN::OpT);
|
|
layernorm->type = OpType_LayerNorm;
|
|
layernorm->main.value = new LayerNormT;
|
|
layernorm->main.type = OpParameter_LayerNorm;
|
|
auto param = layernorm->main.AsLayerNorm();
|
|
param->axis.resize(-axis);
|
|
for (int i=0; i<param->axis.size(); ++i) {
|
|
param->axis[i] = i-(int)(param->axis.size());
|
|
}
|
|
param->epsilon = eps;
|
|
const float* scalePtr = nullptr;
|
|
const float* biasPtr = nullptr;
|
|
if (inputs.size() > 1) {
|
|
scalePtr = inputs[1]->readMap<float>();
|
|
}
|
|
if (nullptr != scalePtr) {
|
|
param->gamma.resize(inputs[1]->getInfo()->size);
|
|
::memcpy(param->gamma.data(), scalePtr, param->gamma.size() * sizeof(float));
|
|
param->beta.resize(inputs[1]->getInfo()->size);
|
|
::memset(param->beta.data(), 0, param->gamma.size() * sizeof(float));
|
|
}
|
|
if (inputs.size() > 2 && nullptr != scalePtr) {
|
|
biasPtr = inputs[2]->readMap<float>();
|
|
}
|
|
if (nullptr != biasPtr) {
|
|
::memcpy(param->beta.data(), biasPtr, param->gamma.size() * sizeof(float));
|
|
}
|
|
auto layerexpr = Expr::create(layernorm.get(), {input});
|
|
auto output = Variable::create(layerexpr);
|
|
if (scalePtr == nullptr) {
|
|
if (inputs.size() > 1) {
|
|
output = output * inputs[1];
|
|
}
|
|
}
|
|
if (biasPtr == nullptr) {
|
|
if (inputs.size() > 2) {
|
|
output = output + inputs[2];
|
|
}
|
|
}
|
|
output->setName(expr->name());
|
|
return output->expr().first;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("BatchNormalization",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxBatchNormTransform));
|
|
OnnxExtraManager::get()->insert("InstanceNormalization",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxInstanceNormalTransform));
|
|
OnnxExtraManager::get()->insert("MeanVarianceNormalization",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxMeanVarianceNormTransform));
|
|
OnnxExtraManager::get()->insert("LpNormalization",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxLpNormTransform));
|
|
OnnxExtraManager::get()->insert("LayerNormalization",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxLayerNormTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|