wehub-resource-sync
55 lines
1.6 KiB
C++
55 lines
1.6 KiB
C++
//
|
|
// TFDense.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/27.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "MNN_generated.h"
|
|
#include "TFExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
class DenseTransform : public TFExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto op = expr->get();
|
|
MNN_ASSERT(nullptr != op->main_as_Extra());
|
|
auto extraAttr = op->main_as_Extra()->attr();
|
|
bool transposeA = false;
|
|
bool transposeB = false;
|
|
bool bias = false;
|
|
if (nullptr != extraAttr) {
|
|
for (int i = 0; i < extraAttr->size(); ++i) {
|
|
auto attr = extraAttr->GetAs<Attribute>(i);
|
|
if ("use_bias" == attr->key()->str()) {
|
|
bias = attr->b();
|
|
continue;
|
|
}
|
|
if ("transpose_b" == attr->key()->str()) {
|
|
transposeB = attr->b();
|
|
continue;
|
|
}
|
|
if ("transpose_a" == attr->key()->str()) {
|
|
transposeA = attr->b();
|
|
continue;
|
|
}
|
|
}
|
|
}
|
|
auto output = _MatMul(inputs[0], inputs[1], transposeA, transposeB);
|
|
if (bias) {
|
|
output = output + inputs[2];
|
|
}
|
|
output->setName(expr->name());
|
|
return output->expr().first;
|
|
}
|
|
};
|
|
static auto gRegister = []() {
|
|
TFExtraManager::get()->insert("Dense", std::shared_ptr<TFExtraManager::Transform>(new DenseTransform));
|
|
return true;
|
|
}();
|
|
} // namespace Express
|
|
} // namespace MNN
|