wehub-resource-sync
94 lines
3.1 KiB
C++
94 lines
3.1 KiB
C++
//
|
|
// TFExtraOp.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/16.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
/*Treat Simple Extra Op*/
|
|
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
#include "MNN_generated.h"
|
|
#include "TFExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
static VARP _Cast(VARP x, DataType src, DataType dst) {
|
|
std::unique_ptr<OpT> castOp(new OpT);
|
|
castOp->type = OpType_Cast;
|
|
castOp->main.value = new CastParamT;
|
|
castOp->main.type = OpParameter_CastParam;
|
|
castOp->main.AsCastParam()->srcT = src;
|
|
castOp->main.AsCastParam()->dstT = dst;
|
|
return Variable::create(Expr::create(castOp.get(), {x}));
|
|
}
|
|
class LogSoftmaxTransform : public TFExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto op = expr->get();
|
|
auto inputs = expr->inputs();
|
|
|
|
int axis = -1;
|
|
if (nullptr != op->main_as_Extra()->attr()) {
|
|
auto n = op->main_as_Extra()->attr()->size();
|
|
for (int i = 0; i < n; ++i) {
|
|
auto attr = op->main_as_Extra()->attr()->GetAs<Attribute>(i);
|
|
if (attr->key()->str() == "axis") {
|
|
axis = attr->i();
|
|
break;
|
|
}
|
|
}
|
|
}
|
|
VARP x = inputs[0];
|
|
VARP max = _ReduceMax(x, {axis}, true);
|
|
VARP sum = _ReduceSum(_Exp(x - max), {axis}, true);
|
|
VARP newVar = x - max - _Log(sum);
|
|
return newVar->expr().first;
|
|
}
|
|
};
|
|
|
|
class LogicalNotTransform : public TFExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
// For bool tensor, the value is 1 or 0
|
|
// 1 - value is equal to logcicalNot
|
|
auto inputs = expr->inputs();
|
|
auto one = _Scalar<int32_t>(1);
|
|
auto newVar = (one - inputs[0]);
|
|
return newVar->expr().first;
|
|
}
|
|
};
|
|
|
|
VARP _BroadcastToForward(VARP a, VARP shape) {
|
|
std::unique_ptr<OpT> op(new OpT);
|
|
op->type = OpType_BroadcastTo;
|
|
op->main.type = OpParameter_Axis;
|
|
auto param = new AxisT;
|
|
param->axis = 1;
|
|
op->main.value = param;
|
|
return (Variable::create(Expr::create(std::move(op), {a, shape})));
|
|
}
|
|
|
|
class SelectTransform : public TFExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
MNN_ASSERT(inputs.size() == 3);
|
|
auto cond = inputs[0];
|
|
auto tvalue = inputs[1];
|
|
auto fvalue = inputs[1];
|
|
auto newVar = _Select(_BroadcastToForward(inputs[0], _Shape(inputs[1])), inputs[1], inputs[2]);
|
|
return newVar->expr().first;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
TFExtraManager::get()->insert("LogicalNot", std::shared_ptr<TFExtraManager::Transform>(new LogicalNotTransform));
|
|
TFExtraManager::get()->insert("LogSoftmax", std::shared_ptr<TFExtraManager::Transform>(new LogSoftmaxTransform));
|
|
TFExtraManager::get()->insert("Select", std::shared_ptr<TFExtraManager::Transform>(new SelectTransform));
|
|
return true;
|
|
}();
|
|
} // namespace Express
|
|
} // namespace MNN
|