wehub-resource-sync
249 lines
9.3 KiB
C++
249 lines
9.3 KiB
C++
//
|
|
// Transformer.cpp
|
|
// MNN
|
|
//
|
|
// Created by MNN on 2019/12/16.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "Transformer.hpp"
|
|
#include "OpConverter.hpp"
|
|
#include "MNN_generated.h"
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
using namespace MNN::Express;
|
|
namespace MNN {
|
|
namespace Train {
|
|
|
|
bool TurnTrainable::onExecute(const std::vector<VARP>& outputs, std::shared_ptr<Parameters> p) {
|
|
auto& trainInfo = mTrainInfo.bnVariables;
|
|
auto exprs = Variable::getExecuteOrder(outputs);
|
|
{
|
|
auto isTraining = _Input({}, NCHW, halide_type_of<int>());
|
|
isTraining->setName("is_training");
|
|
trainInfo["is_training"] = isTraining;
|
|
isTraining = _Cast<float>(isTraining);
|
|
isTraining->setName("is_training_float");
|
|
trainInfo["is_training_float"] = isTraining;
|
|
trainInfo["one_float"] = _Scalar<float>(1.0f);
|
|
trainInfo["bn_momentum"] = _Scalar<float>(mConfig.extraParams["BatchNorm"]["momentum"]->f);
|
|
// Turn convolution be trainable convolution
|
|
for (auto expr : exprs) {
|
|
auto newExpr = OpConverter::convert(expr, mTrainInfo);
|
|
if (newExpr.get() != expr.get()) {
|
|
Expr::replace(expr, newExpr);
|
|
}
|
|
}
|
|
}
|
|
exprs = Variable::getExecuteOrder(outputs);
|
|
auto& noUpdateOps = mConfig.noUpdateOps;
|
|
auto& onlyUpdateOps = mConfig.onlyUpdateOps;
|
|
// Collect Const Variable and turn to Trainable
|
|
for (auto v : exprs) {
|
|
if (v->get() == nullptr && VARP::INPUT != v->inputType()) {
|
|
auto name = v->name();
|
|
auto info = v->outputInfo(0);
|
|
if (halide_type_float != info->type.code) {
|
|
continue;
|
|
}
|
|
|
|
bool update;
|
|
if (!onlyUpdateOps.empty()) {
|
|
update = false;
|
|
for (auto limit : onlyUpdateOps) {
|
|
if (name.find(limit) != std::string::npos) {
|
|
update = true;
|
|
break;
|
|
}
|
|
}
|
|
} else {
|
|
update = true;
|
|
for (auto limit : noUpdateOps) {
|
|
if (name.find(limit) != std::string::npos) {
|
|
update = false;
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
|
|
auto va = Variable::create(v, 0);
|
|
if (update && name != "") {
|
|
va.fix(VARP::TRAINABLE);
|
|
if (name.find("Weight") == std::string::npos && name.find("Bias") == std::string::npos) {
|
|
MNN_PRINT(">>>\ncheck mnn model if const '%s' is a learnable parameter in your original training model, ", name.c_str());
|
|
MNN_PRINT("if not, add it to transformConfig.json NoUpdateOps\n<<<\n");
|
|
va->setName(name + "_Weight");
|
|
va->expr().first->setName(va->name());
|
|
}
|
|
mTrainInfo.trainables.insert(std::make_pair(name, va->name()));
|
|
MNN_PRINT("Add Variable: %s\n", va->name().c_str());
|
|
} else {
|
|
va.fix(VARP::CONSTANT);
|
|
}
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::shared_ptr<Express::Optimizer> Transformer::turnModelToTrainable(TrainConfig config) {
|
|
std::shared_ptr<Express::Optimizer> res;
|
|
res.reset(new TurnTrainable(std::move(config)));
|
|
return res;
|
|
}
|
|
|
|
bool InferOptimizer::onExecute(const std::vector<VARP>& outputs, std::shared_ptr<Parameters> parameters) {
|
|
auto exprs = Variable::getExecuteOrder(outputs);
|
|
|
|
// convert trainable to const
|
|
for (auto& expr : exprs) {
|
|
if (expr->inputs().size() == 0 && expr->inputType() == VARP::InputType::TRAINABLE) {
|
|
auto newConst = Variable::create(expr);
|
|
newConst.fix(VARP::InputType::CONSTANT);
|
|
newConst->setName(expr->name());
|
|
auto newExpr = newConst->expr().first;
|
|
newExpr->setName(expr->name());
|
|
Expr::replace(expr, newExpr);
|
|
}
|
|
}
|
|
|
|
// merge bn after conv into conv
|
|
// convert single bn to scale
|
|
std::set<std::string> bnNames;
|
|
std::string pattern1 = "_MNN_BN_after_conv_first_op";
|
|
std::string pattern2 = "_MNN_single_BN_first_op";
|
|
for (auto& expr : exprs) {
|
|
if (expr->name().find(pattern1) != std::string::npos) {
|
|
std::string bnName = expr->name();
|
|
for (int i = 0; i < pattern1.size(); i++) {
|
|
bnName.pop_back();
|
|
}
|
|
bnNames.insert(bnName);
|
|
}
|
|
if (expr->name().find(pattern2) != std::string::npos) {
|
|
std::string bnName = expr->name();
|
|
for (int i = 0; i < pattern2.size(); i++) {
|
|
bnName.pop_back();
|
|
}
|
|
bnNames.insert(bnName);
|
|
}
|
|
}
|
|
|
|
std::map<std::string, std::map<std::string, EXPRP>> bnInfo;
|
|
for (auto& name : bnNames) {
|
|
for (auto& expr : exprs) {
|
|
auto inputs = expr->inputs();
|
|
if (expr->name() == name) {
|
|
bnInfo[name]["Self"] = expr;
|
|
}
|
|
if (inputs.size() == 0 && expr->name() == name + "_BN_RunningMean_Weight") {
|
|
bnInfo[name]["RunningMean"] = expr;
|
|
}
|
|
if (inputs.size() == 0 && expr->name() == name + "_BN_RunningVariance_Weight") {
|
|
bnInfo[name]["RunningVariance"] = expr;
|
|
}
|
|
if (inputs.size() == 0 && expr->name() == name + "_BN_Gamma_Weight") {
|
|
bnInfo[name]["Gamma"] = expr;
|
|
}
|
|
if (inputs.size() == 0 && expr->name() == name + "_BN_Beta_Bias") {
|
|
bnInfo[name]["Bias"] = expr;
|
|
}
|
|
if (inputs.size() == 0 && expr->name() == name + "_BN_Eps_Weight") {
|
|
bnInfo[name]["Eps"] = expr;
|
|
}
|
|
if (expr->name() == name + pattern1) {
|
|
bnInfo[name]["FirstOpAfterConv"] = expr;
|
|
}
|
|
if (expr->name() == name + pattern2) {
|
|
bnInfo[name]["FirstOpSingleBN"] = expr;
|
|
}
|
|
}
|
|
}
|
|
for (auto& bn : bnInfo) {
|
|
auto bnName = bn.first;
|
|
auto info = bn.second;
|
|
|
|
bool bnAfterConv = false;
|
|
if (info.find("FirstOpAfterConv") != info.end()) {
|
|
bnAfterConv = true;
|
|
}
|
|
|
|
auto rm = _Convert(Variable::create(info["RunningMean"]), NCHW);
|
|
auto rv = _Convert(Variable::create(info["RunningVariance"]), NCHW);
|
|
auto gamma = _Convert(Variable::create(info["Gamma"]), NCHW);
|
|
auto bias = _Convert(Variable::create(info["Bias"]), NCHW);
|
|
auto eps = Variable::create(info["Eps"]);
|
|
|
|
auto s = _Sqrt(rv + eps);
|
|
auto alpha = gamma / s;
|
|
auto beta = bias - rm / s * gamma;
|
|
|
|
if (bnAfterConv) {
|
|
auto firstOp = info["FirstOpAfterConv"];
|
|
auto convExpr = firstOp->inputs()[0]->expr().first;
|
|
if (convExpr->get() == nullptr || convExpr->get()->type() != OpType_Convolution) {
|
|
continue;
|
|
}
|
|
auto convInput = convExpr->inputs()[0];
|
|
auto w = convExpr->inputs()[1];
|
|
auto b = convExpr->inputs()[2];
|
|
|
|
auto nw = w * _Reshape(alpha, {b->getInfo()->dim[0], 1, 1, 1});
|
|
nw.fix(w->expr().first->inputType());
|
|
nw->setName(w->name());
|
|
auto nb = _Reshape(alpha, {b->getInfo()->dim}) * b + _Reshape(beta, b->getInfo()->dim);
|
|
nb.fix(b->expr().first->inputType());
|
|
nb->setName(b->name());
|
|
|
|
std::vector<VARP> newInputs = {convInput, nw, nb};
|
|
auto newConv = Expr::create(convExpr->extra(), std::move(newInputs));
|
|
Expr::replace(info["Self"], newConv);
|
|
} else {
|
|
auto firstOp = info["FirstOpSingleBN"];
|
|
auto inputs = firstOp->inputs();
|
|
std::vector<float> scales, p;
|
|
for (int i = 0; i < beta->getInfo()->size; i++) {
|
|
scales.push_back(alpha->readMap<float>()[i]);
|
|
p.push_back(beta->readMap<float>()[i]);
|
|
}
|
|
auto res = _Scale(inputs[0], beta->getInfo()->size, std::move(scales), std::move(p));
|
|
res->setName(info["Self"]->name());
|
|
Expr::replace(info["Self"], res->expr().first);
|
|
}
|
|
}
|
|
|
|
exprs = Variable::getExecuteOrder(outputs);
|
|
for (auto& iter : exprs) {
|
|
auto op = iter->get();
|
|
if (nullptr == op) {
|
|
continue;
|
|
}
|
|
if (op->type() != OpType_ConvInt8 && op->type() != OpType_DepthwiseConvInt8) {
|
|
continue;
|
|
}
|
|
auto inputExpr = iter->inputs()[0]->expr().first;
|
|
if (inputExpr->get() == nullptr) {
|
|
continue;
|
|
}
|
|
if (inputExpr->get()->type() != OpType_FloatToInt8) {
|
|
continue;
|
|
}
|
|
auto subInputExpr = inputExpr->inputs()[0]->expr().first;
|
|
if (subInputExpr->get() == nullptr) {
|
|
continue;
|
|
}
|
|
if (subInputExpr->get()->type() != OpType_Int8ToFloat) {
|
|
continue;
|
|
}
|
|
//MNN_PRINT("Find direct\n");
|
|
std::vector<VARP> newInputs = subInputExpr->inputs();
|
|
auto newExpr = Expr::create(iter->extra(), std::move(newInputs));
|
|
Expr::replace(iter, newExpr);
|
|
}
|
|
return true;
|
|
}
|
|
|
|
std::shared_ptr<Express::Optimizer> Transformer::turnModelToInfer() {
|
|
return std::shared_ptr<Optimizer>(new InferOptimizer);
|
|
}
|
|
} // namespace Train
|
|
} // namespace MNN
|