wehub-resource-sync
88 lines
2.9 KiB
C++
88 lines
2.9 KiB
C++
//
|
|
// OnnxExtraManager.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/29.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "OnnxExtraManager.hpp"
|
|
#include "OpCount.hpp"
|
|
#include "MNN_generated.h"
|
|
namespace MNN {
|
|
namespace Express {
|
|
VARP OnnxExtraManager::_ReshapeF(VARP x, VARP shape, int format) {
|
|
MNN_ASSERT(nullptr != x);
|
|
std::unique_ptr<OpT> reshape(new OpT);
|
|
reshape->type = OpType_Reshape;
|
|
reshape->main.type = OpParameter_Reshape;
|
|
reshape->main.value = new ReshapeT;
|
|
reshape->main.AsReshape()->dimType = (MNN_DATA_FORMAT)format;
|
|
return (Variable::create(Expr::create(reshape.get(), {x, shape})));
|
|
}
|
|
std::shared_ptr<OnnxExtraManager> OnnxExtraManager::get() {
|
|
static std::shared_ptr<OnnxExtraManager> gInstance;
|
|
if (nullptr == gInstance) {
|
|
gInstance.reset(new OnnxExtraManager);
|
|
}
|
|
return gInstance;
|
|
}
|
|
|
|
void OnnxExtraManager::insert(const std::string& name, std::shared_ptr<Transform> transform) {
|
|
mTransform.insert(std::make_pair(name, transform));
|
|
OpCount::get()->insertOp("ONNX", name);
|
|
}
|
|
std::shared_ptr<OnnxExtraManager::Transform> OnnxExtraManager::find(const std::string& name) const {
|
|
auto iter = mTransform.find(name);
|
|
if (iter == mTransform.end()) {
|
|
return nullptr;
|
|
}
|
|
return iter->second;
|
|
}
|
|
|
|
static auto gRegister = []() {
|
|
auto extra = OnnxExtraManager::get();
|
|
auto judge = [extra](EXPRP expr) {
|
|
auto op = expr->get();
|
|
if (nullptr == op) {
|
|
return false;
|
|
}
|
|
if (op->type() != OpType_Extra) {
|
|
return false;
|
|
}
|
|
auto engine = op->main_as_Extra()->engine()->str();
|
|
if (engine != "ONNX") {
|
|
return false;
|
|
}
|
|
auto type = op->main_as_Extra()->type()->str();
|
|
if (extra->find(type) == nullptr) {
|
|
return false;
|
|
}
|
|
return true;
|
|
};
|
|
auto modify = [extra](EXPRP expr) {
|
|
auto op = expr->get();
|
|
MNN_ASSERT(op->type() == OpType_Extra);
|
|
auto type = op->main_as_Extra()->type()->str();
|
|
auto transformer = extra->find(type);
|
|
MNN_ASSERT(nullptr != transformer);
|
|
auto newExpr = transformer->onExecute(expr);
|
|
if (nullptr == newExpr) {
|
|
MNN_ERROR("Convert Onnx's Op %s , type = %s, failed, may be some node is not const\n", expr->name().c_str(),
|
|
type.c_str());
|
|
return false;
|
|
}
|
|
if (expr->outputSize() == newExpr->outputSize()) {
|
|
for (int i=0; i<expr->outputSize(); ++i) {
|
|
Variable::create(newExpr, i)->setName(expr->outputName(i));
|
|
}
|
|
}
|
|
Expr::replace(expr, newExpr);
|
|
return true;
|
|
};
|
|
TemplateMerge::getInstance("OnnxExtra").insertTemplate("OnnxExtraManager", judge, modify);
|
|
return true;
|
|
}();
|
|
} // namespace Express
|
|
} // namespace MNN
|