wehub-resource-sync
112 lines
4.0 KiB
C++
112 lines
4.0 KiB
C++
//
|
|
// OnnxSlice.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/09/16.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <climits>
|
|
#include <numeric>
|
|
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
// Returns true if all elements in the array are 1, otherwise returns false.
|
|
static bool IsAllOne(const std::vector<int> &array) {
|
|
for (const int &i : array) {
|
|
if (i != 1)
|
|
return false;
|
|
}
|
|
return array.size() ? true : false;
|
|
}
|
|
|
|
class OnnxSliceTransform : public OnnxExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto op = expr->get();
|
|
MNN_ASSERT(op->type() == OpType_Extra);
|
|
|
|
auto type = op->main_as_Extra()->type()->str();
|
|
auto inputs = expr->inputs();
|
|
auto input = inputs[0];
|
|
auto attrs = op->main_as_Extra()->attr();
|
|
if (inputs.size() == 1 && nullptr == attrs) {
|
|
MNN_PRINT("Attrs of Slice in ONNX must not be null when inputs.size == 1\n");
|
|
return nullptr;
|
|
}
|
|
VARP startVar;
|
|
VARP endVar;
|
|
VARP axisVar;
|
|
VARP strideVar;
|
|
if (nullptr != attrs) {
|
|
// Copy from attribute
|
|
std::vector<int> starts, ends, axes, strides;
|
|
auto copyFunction = [](std::vector<int> &dst, const MNN::Attribute *attr) {
|
|
MNN_ASSERT(nullptr != attr->list());
|
|
MNN_ASSERT(nullptr != attr->list()->i());
|
|
dst.resize(attr->list()->i()->size());
|
|
::memcpy(dst.data(), attr->list()->i()->data(), dst.size() * sizeof(int));
|
|
};
|
|
for (int i = 0; i < attrs->size(); ++i) {
|
|
auto attr = attrs->GetAs<Attribute>(i);
|
|
if (attr->key()->str() == "axes") {
|
|
copyFunction(axes, attr);
|
|
axisVar = _Const(axes.data(), {(int)axes.size()}, NCHW, halide_type_of<int>());
|
|
} else if (attr->key()->str() == "ends") {
|
|
copyFunction(ends, attr);
|
|
endVar = _Const(ends.data(), {(int)ends.size()}, NCHW, halide_type_of<int>());
|
|
} else if (attr->key()->str() == "starts") {
|
|
copyFunction(starts, attr);
|
|
startVar = _Const(starts.data(), {(int)starts.size()}, NCHW, halide_type_of<int>());
|
|
} else if (attr->key()->str() == "steps") {
|
|
copyFunction(strides, attr);
|
|
strideVar = _Const(strides.data(), {(int)strides.size()}, NCHW, halide_type_of<int>());
|
|
}
|
|
}
|
|
}
|
|
{
|
|
// If has input, use input instead of attribute
|
|
if (inputs.size() > 1) {
|
|
startVar = inputs[1];
|
|
}
|
|
if (inputs.size() > 2) {
|
|
endVar = inputs[2];
|
|
}
|
|
if (inputs.size() > 3) {
|
|
axisVar = inputs[3];
|
|
}
|
|
if (inputs.size() > 4) {
|
|
strideVar = inputs[4];
|
|
}
|
|
}
|
|
|
|
std::unique_ptr<MNN::OpT> sliceOp(new OpT);
|
|
sliceOp->name = op->name()->str();
|
|
|
|
sliceOp->type = OpType_StridedSlice;
|
|
sliceOp->main.type = OpParameter_StridedSliceParam;
|
|
auto param = new StridedSliceParamT;
|
|
param->fromType = 1;
|
|
sliceOp->main.value = param;
|
|
if(nullptr != axisVar && nullptr != strideVar) {
|
|
return Expr::create(sliceOp.get(), {input, startVar, endVar, axisVar, strideVar}, expr->outputSize());
|
|
}
|
|
if(nullptr != axisVar) {
|
|
return Expr::create(sliceOp.get(), {input, startVar, endVar, axisVar}, expr->outputSize());
|
|
}
|
|
return Expr::create(sliceOp.get(), {input, startVar, endVar}, expr->outputSize());
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("Slice", std::shared_ptr<OnnxExtraManager::Transform>(new OnnxSliceTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|