wehub-resource-sync
86 lines
3.2 KiB
C++
86 lines
3.2 KiB
C++
//
|
|
// OnnxConstantOfShape.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2019/10/31.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
template <typename T>
|
|
static void ResizeAndCopyData(std::vector<uint8_t> *vec_data, const T *src, const std::vector<int> &shape) {
|
|
int64_t count = sizeof(T);
|
|
for (const auto &d : shape) {
|
|
count *= d;
|
|
}
|
|
MNN_ASSERT(count > 0);
|
|
|
|
vec_data->resize(count);
|
|
memcpy(vec_data->data(), src, count);
|
|
}
|
|
|
|
class OnnxConstantOfShapeTransform : public OnnxExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
auto inputs = expr->inputs();
|
|
MNN_THROW_CHECK(1 == inputs.size(), "Onnx ConstantOfShape should have one input!");
|
|
|
|
std::unique_ptr<OpT> mnnFill(new OpT);
|
|
mnnFill->name = expr->name();
|
|
mnnFill->defaultDimentionFormat = MNN_DATA_FORMAT_NCHW;
|
|
mnnFill->type = OpType_Fill;
|
|
mnnFill->main.type = OpParameter_NONE;
|
|
mnnFill->main.value = nullptr;
|
|
|
|
std::vector<uint8_t> tensor_data;
|
|
std::vector<int> tensor_shape;
|
|
halide_type_t data_type = halide_type_of<float>();
|
|
auto extraParam = expr->get()->main_as_Extra();
|
|
VARP const_shape;
|
|
if (extraParam->attr() != nullptr) {
|
|
for (int i = 0; i < extraParam->attr()->size(); ++i) {
|
|
auto attr = extraParam->attr()->GetAs<Attribute>(i);
|
|
const auto &key = attr->key()->str();
|
|
if (key == "value") {
|
|
auto blob = attr->tensor();
|
|
// Process tensor shape.
|
|
tensor_shape.resize(blob->dims()->size());
|
|
for (int j = 0; j < tensor_shape.size(); ++j) {
|
|
tensor_shape[j] = blob->dims()->Get(j);
|
|
}
|
|
|
|
if (blob->dataType() == DataType_DT_INT32) {
|
|
data_type = halide_type_of<int32_t>();
|
|
ResizeAndCopyData<int>(&tensor_data, blob->int32s()->data(), tensor_shape);
|
|
} else if (blob->dataType() == DataType_DT_FLOAT) {
|
|
data_type = halide_type_of<float>();
|
|
ResizeAndCopyData<float>(&tensor_data, blob->float32s()->data(), tensor_shape);
|
|
} else {
|
|
MNN_ERROR("Not support data type.");
|
|
}
|
|
break; // Break out of the loop if value has been processed.
|
|
}
|
|
}
|
|
const_shape = _Const(static_cast<const void *>(tensor_data.data()), tensor_shape, NCHW, data_type);
|
|
} else {
|
|
// https://github.com/onnx/onnx/blob/master/docs/Operators.md#attributes-12
|
|
const_shape = _Scalar<float>(0.0f);
|
|
}
|
|
return Expr::create(mnnFill.get(), {inputs[0], const_shape});
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("ConstantOfShape",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxConstantOfShapeTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|