wehub-resource-sync
53 lines
1.9 KiB
C++
53 lines
1.9 KiB
C++
//
|
|
// ConvTranposeTflite.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/27.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "MNN_generated.h"
|
|
#include "../../tflite/liteOpConverter.hpp"
|
|
#include "TFliteExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
/*See CustomTflite.cpp for detail attribute*/
|
|
class ConvTranposeTflite : public TFliteExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
auto weight = inputs[1];
|
|
auto bias = inputs[2];
|
|
weight = _Transpose(weight, {3, 0, 1, 2});
|
|
auto weightInfo = weight->getInfo();
|
|
auto biasInfo = bias->getInfo();
|
|
|
|
auto extra = expr->get()->main_as_Extra();
|
|
std::unique_ptr<MNN::OpT> deconvOp(flatbuffers::GetRoot<MNN::Op>(extra->info()->data())->UnPack());
|
|
auto weightPtr = weight->readMap<float>();
|
|
auto biasPtr = bias->readMap<float>();
|
|
EXPRP newExpr;
|
|
if (nullptr == weightPtr || nullptr == biasPtr) {
|
|
newExpr = Expr::create(deconvOp.get(), {inputs[0], weight, bias});
|
|
} else {
|
|
auto conv = deconvOp->main.AsConvolution2D();
|
|
conv->weight.resize(weightInfo->size);
|
|
::memcpy(conv->weight.data(), weightPtr, weightInfo->size * sizeof(float));
|
|
conv->bias.resize(biasInfo->size);
|
|
::memcpy(conv->bias.data(), biasPtr, biasInfo->size * sizeof(float));
|
|
newExpr = Expr::create(deconvOp.get(), {inputs[0]});
|
|
}
|
|
auto newOutput = Variable::create(newExpr);
|
|
newOutput->setName(expr->name());
|
|
return newOutput->expr().first;
|
|
}
|
|
};
|
|
static auto gRegister = []() {
|
|
TFliteExtraManager::get()->insert("Convolution2DTransposeBias", std::shared_ptr<TFliteExtraManager::Transform>(new ConvTranposeTflite));
|
|
return true;
|
|
}();
|
|
} // namespace Express
|
|
} // namespace MNN
|