wehub-resource-sync
92 lines
3.4 KiB
C++
92 lines
3.4 KiB
C++
//
|
|
// TFConvolution3DMerge.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/12/03.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <algorithm>
|
|
#include "MNN_generated.h"
|
|
#include "TFExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class Convolution3DTransform : public TFExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto op = expr->get();
|
|
auto inputs = expr->inputs();
|
|
auto weight = inputs[1];
|
|
auto weightInfo = weight->getInfo();
|
|
auto weightTensorData = weight->readMap<float>();
|
|
if (nullptr == weightInfo || nullptr == weightTensorData) {
|
|
MNN_ERROR("For %s Convolution3D weight is not const\n", expr->name().c_str());
|
|
return nullptr;
|
|
}
|
|
|
|
std::unique_ptr<Convolution3DT> conv3d(new MNN::Convolution3DT);
|
|
int depth = weightInfo->dim[0];
|
|
int kh = weightInfo->dim[1];
|
|
int kw = weightInfo->dim[2];
|
|
int num_input = weightInfo->dim[3];
|
|
int num_output = weightInfo->dim[4];
|
|
weight = _Transpose(weight, {4, 3, 0, 1, 2});
|
|
weightInfo = weight->getInfo();
|
|
weightTensorData = weight->readMap<float>();
|
|
conv3d->bias.resize(num_output);
|
|
std::fill(conv3d->bias.begin(), conv3d->bias.end(), 0.0f);
|
|
|
|
conv3d->weight.resize(weightInfo->size);
|
|
::memcpy(conv3d->weight.data(), weightTensorData, weightInfo->size * sizeof(float));
|
|
conv3d->common.reset(new MNN::Convolution3DCommonT);
|
|
auto common = conv3d->common.get();
|
|
|
|
common->relu = common->relu6 = false;
|
|
common->outputCount = num_output;
|
|
common->inputCount = num_input;
|
|
common->kernels = std::vector<int>({depth, kh, kw});
|
|
|
|
auto extra = op->main_as_Extra();
|
|
if (extra == nullptr || extra->attr() == nullptr) {
|
|
return nullptr;
|
|
}
|
|
for (int i = 0; i < extra->attr()->size(); ++i) {
|
|
auto attr = extra->attr()->GetAs<Attribute>(i);
|
|
const auto key = attr->key()->str();
|
|
if (key == "dilations" || key == "rates") {
|
|
auto values = attr->list()->i()->data();
|
|
common->dilates = std::vector<int>({values[1], values[2], values[3]});
|
|
} else if (key == "strides") {
|
|
auto values = attr->list()->i()->data();
|
|
common->strides = std::vector<int>({values[1], values[2], values[3]});
|
|
} else if (key == "padding") {
|
|
common->padMode = MNN::PadMode_SAME;
|
|
auto paddingType = attr->s()->str();
|
|
if (paddingType == "VALID") {
|
|
common->padMode = MNN::PadMode_VALID;
|
|
common->pads = std::vector<int>({0, 0, 0});
|
|
}
|
|
}
|
|
}
|
|
|
|
std::unique_ptr<OpT> newOp(new OpT);
|
|
newOp->name = expr->name();
|
|
newOp->type = OpType_Convolution3D;
|
|
newOp->main.type = OpParameter_Convolution3D;
|
|
newOp->main.value = conv3d.release();
|
|
|
|
auto newExpr = Expr::create(newOp.get(), {inputs[0]}, 1);
|
|
return newExpr;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
TFExtraManager::get()->insert("Conv3D", std::shared_ptr<TFExtraManager::Transform>(new Convolution3DTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|