wehub-resource-sync
78 lines
2.9 KiB
C++
78 lines
2.9 KiB
C++
//
|
|
// OnnxUnPool.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2020/01/21.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
/*
|
|
MaxUnPool implemention is same as onnxruntime / pytorch
|
|
Note: test case in onnx's docs is wrong. https://github.com/onnx/onnx/issues/2398
|
|
*/
|
|
class OnnxUnPoolTransform : public OnnxExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
INTS kernels, pads, strides;
|
|
auto attrs = expr->get()->main_as_Extra()->attr();
|
|
auto extractVec = [](INTS& vec, const Attribute* attr) {
|
|
vec.assign(attr->list()->i()->begin(), attr->list()->i()->end());
|
|
};
|
|
for (int i = 0; i < attrs->size(); ++i) {
|
|
auto attr = attrs->GetAs<Attribute>(i);
|
|
const auto& key = attr->key()->str();
|
|
if (key == "kernel_shape") {
|
|
extractVec(kernels, attr);
|
|
} else if (key == "pads") {
|
|
extractVec(pads, attr);
|
|
} else if (key == "strides") {
|
|
extractVec(strides, attr);
|
|
}
|
|
}
|
|
auto inputs = expr->inputs();
|
|
VARP outShape;
|
|
if (inputs.size() == 3) {
|
|
outShape = inputs[2];
|
|
} else {
|
|
int len = kernels.size();
|
|
MNN_THROW_CHECK(strides.size() == 0 || strides.size() == len, "Invalid strides attr");
|
|
MNN_THROW_CHECK(pads.size() == 0 || pads.size() == len * 2, "Invalid pads attr");
|
|
auto kernelV = _Const(kernels.data(), {len}, NCHW, halide_type_of<int>());
|
|
auto oneV = _Scalar<int>(1), zeroV = _Scalar<int>(0);
|
|
auto strideV = oneV, padV = zeroV;
|
|
if (strides.size() != 0) {
|
|
strideV = _Const(strides.data(), {len}, NCHW, halide_type_of<int>());
|
|
}
|
|
if (pads.size() != 0) {
|
|
INTS newPads;
|
|
for (int i = 0; i < len; ++i) {
|
|
newPads.push_back(pads[i] + pads[i + len]);
|
|
}
|
|
padV = _Const(newPads.data(), {len}, NCHW, halide_type_of<int>());
|
|
}
|
|
auto inShape = _Shape(inputs[0], NCHW), twoV = _Unsqueeze(_Scalar<int>(2), {0});
|
|
outShape = _Slice(inShape, twoV, _Unsqueeze(_Scalar<int>(len), {0}));
|
|
outShape = kernelV + (outShape - oneV) * strideV - padV;
|
|
outShape = _Concat({_Slice(inShape, _Unsqueeze(zeroV, {0}), twoV), outShape}, 0);
|
|
}
|
|
auto res = _ScatterNd(_Reshape(inputs[1], {-1, 1}), _Reshape(inputs[0], {-1}), _Unsqueeze(_ReduceProd(outShape), {0}));
|
|
res = _Reshape(res, outShape);
|
|
res->setName(expr->outputName(0));
|
|
return res->expr().first;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("MaxUnpool", std::shared_ptr<OnnxExtraManager::Transform>(new OnnxUnPoolTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|