91 lines
3.1 KiB
C++
91 lines
3.1 KiB
C++
//
|
|
// ConvBiasAdd.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/16.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../TemplateMerge.hpp"
|
|
#include "MNN_generated.h"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
static auto gRegister = []() {
|
|
auto compare = [](EXPRP expr) {
|
|
if (nullptr == expr->get()) {
|
|
return false;
|
|
}
|
|
if (expr->get()->type() != OpType_BinaryOp) {
|
|
return false;
|
|
}
|
|
if (expr->get()->main_as_BinaryOp()->opType() != BinaryOpOperation_ADD) {
|
|
return false;
|
|
}
|
|
auto inputs = expr->inputs();
|
|
auto inputExpr = inputs[0]->expr().first;
|
|
if (nullptr == inputExpr->get()) {
|
|
return false;
|
|
}
|
|
if (inputExpr->get()->type() == OpType_Reshape) {
|
|
inputExpr = inputExpr->inputs()[0]->expr().first;
|
|
}
|
|
if (!inputExpr->get() || inputExpr->get()->main_type() != OpParameter_Convolution2D || inputExpr->outputs().size() != 1) {
|
|
return false;
|
|
}
|
|
if (inputExpr->inputs().size() > 1) {
|
|
return false;
|
|
}
|
|
// Merge into convolution
|
|
auto biasVar = inputs[1];
|
|
auto biasInfo = biasVar->getInfo();
|
|
auto biasPtr = biasVar->readMap<float>();
|
|
if (nullptr == biasInfo || nullptr == biasPtr) {
|
|
return false;
|
|
}
|
|
auto paraent = inputExpr->inputs();
|
|
auto outputCount = inputExpr->get()->main_as_Convolution2D()->common()->outputCount();
|
|
if (biasInfo->size != outputCount) {
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
auto modify = [](EXPRP expr) {
|
|
auto inputs = expr->inputs();
|
|
auto inputExpr = inputs[0]->expr().first;
|
|
auto biasVar = inputs[1];
|
|
auto biasInfo = biasVar->getInfo();
|
|
auto biasPtr = biasVar->readMap<float>();
|
|
EXPRP reshapeExpr = nullptr;
|
|
if (inputExpr->get()->type() == OpType_Reshape) {
|
|
reshapeExpr = inputExpr;
|
|
inputExpr = inputExpr->inputs()[0]->expr().first;
|
|
}
|
|
std::unique_ptr<OpT> convOp(inputExpr->get()->UnPack());
|
|
auto& biasData = convOp->main.AsConvolution2D()->bias;
|
|
MNN_ASSERT(biasInfo->size == biasData.size());
|
|
for (int i = 0; i < biasData.size(); ++i) {
|
|
biasData[i] += biasPtr[i];
|
|
}
|
|
auto newExpr = Expr::create(convOp.get(), inputExpr->inputs());
|
|
if (reshapeExpr) {
|
|
auto convVar = Variable::create(newExpr);
|
|
auto inputs = reshapeExpr->inputs();
|
|
std::vector<VARP> newInputs(inputs.size());
|
|
newInputs[0] = convVar;
|
|
if (inputs.size() == 2) {
|
|
newInputs[1] = inputs[1];
|
|
}
|
|
std::unique_ptr<OpT> reshapeOp(reshapeExpr->get()->UnPack());
|
|
newExpr = Expr::create(reshapeOp.get(), newInputs);
|
|
}
|
|
newExpr->setName(expr->name());
|
|
Expr::replace(expr, newExpr);
|
|
return true;
|
|
};
|
|
TemplateMerge::getInstance("Merge").insertTemplate("ConvBiasAdd", compare, modify);
|
|
return true;
|
|
}();
|
|
}
|
|
} // namespace MNN
|