wehub-resource-sync
69 lines
2.1 KiB
C++
69 lines
2.1 KiB
C++
//
|
|
// OnnxSoftmax.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2023/03/03.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <MNN/expr/Expr.hpp>
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
VARP _ConstInt(int x) {
|
|
return _Const(&x, {1}, NHWC, halide_type_of<int>());
|
|
}
|
|
|
|
class OnnxSoftmaxTransform : public OnnxExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto op = expr->get();
|
|
MNN_ASSERT(op->type() == OpType_Extra);
|
|
int axis = -1, opsetVersion = 13;
|
|
auto attrs = op->main_as_Extra()->attr();
|
|
if (nullptr != attrs) {
|
|
for (int i = 0; i < attrs->size(); ++i) {
|
|
auto attr = attrs->GetAs<Attribute>(i);
|
|
if (attr->key()->str() == "axis") {
|
|
axis = attr->i();
|
|
}
|
|
if (attr->key()->str() == "onnx_opset_version") {
|
|
opsetVersion = attr->i();
|
|
}
|
|
}
|
|
}
|
|
auto input = expr->inputs()[0];
|
|
if (opsetVersion >= 13 || axis == -1) {
|
|
auto newExpr = _Softmax(input, axis)->expr().first;
|
|
newExpr->setName(expr->name());
|
|
return newExpr;
|
|
}
|
|
auto shape = _Shape(input, true);
|
|
{
|
|
std::unique_ptr<MNN::OpT> flat(new OpT);
|
|
flat->type = OpType_Flatten;
|
|
flat->main.type = OpParameter_Flatten;
|
|
flat->main.value = new FlattenT;
|
|
flat->main.AsFlatten()->endAxis = 0;
|
|
flat->main.AsFlatten()->axis = axis;
|
|
input = Variable::create(Expr::create(flat.get(), {input}));
|
|
}
|
|
auto output = _Softmax(input, -1);
|
|
auto newExpr = _Reshape(output, shape)->expr().first;
|
|
newExpr->setName(expr->name());
|
|
return newExpr;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("Softmax",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxSoftmaxTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|