Files
2026-07-13 13:33:03 +08:00

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