wehub-resource-sync
92 lines
3.1 KiB
C++
92 lines
3.1 KiB
C++
//
|
|
// DepthwiseConvWeightMerge.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/04/19.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../TemplateMerge.hpp"
|
|
#include "MNN/expr/MathOp.hpp"
|
|
#include "MNN/expr/NeuralNetWorkOp.hpp"
|
|
#include "MNN_generated.h"
|
|
#include "config.hpp"
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
static auto gRegister = []() {
|
|
|
|
auto match = [](EXPRP expr) {
|
|
auto config = Global<modelConfig>::Get();
|
|
auto modelType = config->model;
|
|
if (modelConfig::TFLITE != modelType) {
|
|
return false;
|
|
}
|
|
if (nullptr == expr->get()) {
|
|
return false;
|
|
}
|
|
if (expr->get()->type() != OpType_ConvolutionDepthwise && expr->get()->type() != OpType_Convolution) {
|
|
return false;
|
|
}
|
|
// 1. input, 2. weight, 3. bias
|
|
auto inputs = expr->inputs();
|
|
if (inputs.size() < 2) {
|
|
return false;
|
|
}
|
|
if (inputs.size() >= 2) {
|
|
auto weightVar = inputs[1];
|
|
auto weightInfo = weightVar->getInfo();
|
|
auto weightPtr = weightVar->readMap<float>();
|
|
if (nullptr == weightInfo || nullptr == weightPtr) {
|
|
return false;
|
|
}
|
|
}
|
|
if (inputs.size() == 3) {
|
|
auto biasVar = inputs[1];
|
|
auto biasInfo = biasVar->getInfo();
|
|
auto biasPtr = biasVar->readMap<float>();
|
|
if (nullptr == biasInfo || nullptr == biasPtr) {
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
};
|
|
|
|
auto transform = [](EXPRP expr) {
|
|
std::unique_ptr<OpT> convOp(expr->get()->UnPack());
|
|
auto inputs = expr->inputs();
|
|
if (inputs.size() >= 2) {
|
|
auto weightVar = inputs[1];
|
|
if (expr->get()->type() == OpType_ConvolutionDepthwise) {
|
|
weightVar = _Transpose(weightVar, {3, 0, 1, 2});
|
|
} else if (expr->get()->type() == OpType_Convolution) {
|
|
weightVar = _Transpose(weightVar, {0, 3, 1, 2});
|
|
}
|
|
auto weightInfo = weightVar->getInfo();
|
|
auto weightPtr = weightVar->readMap<float>();
|
|
auto& weightData = convOp->main.AsConvolution2D()->weight;
|
|
weightData.resize(weightInfo->size);
|
|
memcpy(weightData.data(), weightPtr, weightInfo->size * sizeof(float));
|
|
}
|
|
auto& biasData = convOp->main.AsConvolution2D()->bias;
|
|
biasData.resize(convOp->main.AsConvolution2D()->common->outputCount);
|
|
::memset(biasData.data(), 0, sizeof(float) * biasData.size());
|
|
if (inputs.size() == 3) {
|
|
auto biasVar = inputs[2];
|
|
auto biasInfo = biasVar->getInfo();
|
|
auto biasPtr = biasVar->readMap<float>();
|
|
memcpy(biasData.data(), biasPtr, biasInfo->size * sizeof(float));
|
|
}
|
|
auto newExpr = Expr::create(convOp.get(), {inputs[0]});
|
|
newExpr->setName(expr->name());
|
|
Expr::replace(expr, newExpr);
|
|
return true;
|
|
};
|
|
|
|
TemplateMerge::getInstance("Merge").insertTemplate("DepthwiseConvWeightMerge", match, transform);
|
|
return true;
|
|
}();
|
|
|
|
}
|
|
} // namespace MNN
|