wehub-resource-sync
97 lines
3.0 KiB
C++
97 lines
3.0 KiB
C++
//
|
|
// EliminateSqueezeExpandDims.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/12/03.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "../TemplateMerge.hpp"
|
|
#include "MNN/expr/ExprCreator.hpp"
|
|
#include "MNN_generated.h"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class EliminateSqueezeExpandDims {
|
|
public:
|
|
EliminateSqueezeExpandDims();
|
|
};
|
|
|
|
EliminateSqueezeExpandDims::EliminateSqueezeExpandDims() {
|
|
auto match = [this](EXPRP expr) -> bool {
|
|
if (!expr->get()) {
|
|
return false;
|
|
}
|
|
if ((expr->get()->type() != OpType_Squeeze) && (expr->get()->type() != OpType_ExpandDims)) {
|
|
return false;
|
|
}
|
|
|
|
VARP input = expr->inputs().at(0);
|
|
const Op* inputOp = input->expr().first->get();
|
|
if (inputOp == nullptr) {
|
|
return false;
|
|
}
|
|
|
|
if (input->expr().first->outputSize() != 1) {
|
|
return false;
|
|
}
|
|
|
|
if (expr->get()->type() == OpType_Squeeze) {
|
|
if (inputOp->type() != OpType_ExpandDims) {
|
|
return false;
|
|
}
|
|
auto squeezeDims = expr->get()->main_as_SqueezeParam()->squeezeDims();
|
|
int expandDim = inputOp->main_as_ExpandDims()->axis();
|
|
if (squeezeDims->size() != 1) { // squeeze can apply to multi-dimension, but expand_dims can only have single axis value
|
|
return false;
|
|
}
|
|
if (expandDim != squeezeDims->data()[0]) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
if (expr->get()->type() == OpType_ExpandDims) {
|
|
if (inputOp->type() != OpType_Squeeze) {
|
|
return false;
|
|
}
|
|
auto squeezeDims = inputOp->main_as_SqueezeParam()->squeezeDims();
|
|
int expandDim = expr->get()->main_as_ExpandDims()->axis();
|
|
if (squeezeDims->size() != 1) { // squeeze can apply to multi-dimension, but expand_dims can only have single axis value
|
|
return false;
|
|
}
|
|
if (expandDim != squeezeDims->data()[0]) {
|
|
return false;
|
|
}
|
|
}
|
|
|
|
return true;
|
|
};
|
|
|
|
auto fold = [this](EXPRP expr) -> bool {
|
|
VARP input = expr->inputs().at(0);
|
|
input = input->expr().first->inputs().at(0);
|
|
|
|
auto* identity = new MNN::ExtraT;
|
|
identity->type = "Identity";
|
|
identity->engine = "Tensorflow";
|
|
std::unique_ptr<MNN::OpT> identity_op(new MNN::OpT);
|
|
identity_op->name = expr->name();
|
|
identity_op->type = OpType_Extra;
|
|
identity_op->main.type = OpParameter_Extra;
|
|
identity_op->main.value = identity;
|
|
|
|
EXPRP identity_expr = Expr::create(identity_op.get(), {input});
|
|
Expr::replace(expr, identity_expr);
|
|
|
|
return true /*modified*/;
|
|
};
|
|
|
|
TemplateMerge::getInstance("Merge").insertTemplate("EliminateSqueezeExpandDims", match, fold);
|
|
}
|
|
|
|
static EliminateSqueezeExpandDims g_eliminate_squeeze_expand_dims;
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|