wehub-resource-sync
106 lines
4.4 KiB
C++
106 lines
4.4 KiB
C++
//
|
|
// OnnxRandomUniform.cpp
|
|
// MNNConverter
|
|
//
|
|
// Created by MNN on 2021/04/20.
|
|
// Copyright © 2018, Alibaba Group Holding Limited
|
|
//
|
|
|
|
#include <MNN/expr/ExprCreator.hpp>
|
|
#include "MNN_generated.h"
|
|
#include "OnnxExtraManager.hpp"
|
|
|
|
namespace MNN {
|
|
namespace Express {
|
|
|
|
class OnnxRandomUniformTransform : public OnnxExtraManager::Transform {
|
|
public:
|
|
virtual EXPRP onExecute(EXPRP expr) const override {
|
|
static bool gInit = false;
|
|
if (!gInit) {
|
|
MNN_PRINT("The model has random OP: %s, can't check result with onnxruntime\n", expr->name().c_str());
|
|
gInit = true;
|
|
}
|
|
auto op = expr->get();
|
|
auto info = op->main_as_Extra();
|
|
std::unique_ptr<OpT> randomUniform(new OpT);
|
|
randomUniform->name = expr->name();
|
|
if (info->type()->str() == "RandomUniform" || info->type()->str() == "RandomUniformLike") {
|
|
randomUniform->type = OpType_RandomUniform;
|
|
} else {
|
|
randomUniform->type = OpType_RandomNormal;
|
|
}
|
|
randomUniform->main.type = OpParameter_RandomUniform;
|
|
auto param = new RandomUniformT;
|
|
randomUniform->main.value = param;
|
|
std::vector<int> outoutShape;
|
|
bool hasShape = false;
|
|
bool hasType = false;
|
|
if (nullptr != info->attr()) {
|
|
for (int i = 0; i < info->attr()->size(); ++i) {
|
|
const auto attr = info->attr()->GetAs<Attribute>(i);
|
|
const auto attributeName = attr->key()->str();
|
|
if (attributeName == "shape") {
|
|
if (nullptr != attr->list() && nullptr != attr->list()->i()) {
|
|
outoutShape.resize(attr->list()->i()->size());
|
|
::memcpy(outoutShape.data(), attr->list()->i()->data(), outoutShape.size() * sizeof(int));
|
|
}
|
|
hasShape = true;
|
|
} else if (attributeName == "low" || attributeName == "mean") {
|
|
param->low = attr->f();
|
|
} else if (attributeName == "high" || attributeName == "scale") {
|
|
param->high = attr->f();
|
|
} else if (attributeName == "seed") {
|
|
param->seed = attr->i();
|
|
} else if (attributeName == "dtype") {
|
|
param->type = static_cast<MNN::DataType>(attr->i());
|
|
hasType = true;
|
|
}
|
|
}
|
|
}
|
|
EXPRP newExpr;
|
|
if (hasShape) {
|
|
auto const_shape = _Const(static_cast<const void *>(outoutShape.data()), { static_cast<int>(outoutShape.size()) }, NCHW, halide_type_of<int>());
|
|
newExpr = Expr::create(randomUniform.get(), {const_shape});
|
|
} else {
|
|
if (!hasType) {
|
|
auto info = expr->inputs()[0]->getInfo();
|
|
if (nullptr == info) {
|
|
return nullptr;
|
|
}
|
|
if (info->type.code == halide_type_float) {
|
|
param->type = DataType_DT_FLOAT;
|
|
} else if (info->type.code == halide_type_int) {
|
|
if (info->type.bytes() == 1) {
|
|
param->type = DataType_DT_INT8;
|
|
} else {
|
|
param->type = DataType_DT_INT32;
|
|
}
|
|
} else if (info->type.code == halide_type_uint) {
|
|
if (info->type.bytes() == 1) {
|
|
param->type = DataType_DT_UINT8;
|
|
}
|
|
}
|
|
}
|
|
newExpr = Expr::create(randomUniform.get(), {_Shape(expr->inputs()[0], NCHW)});
|
|
}
|
|
newExpr->setName(expr->name());
|
|
return newExpr;
|
|
}
|
|
};
|
|
|
|
static auto gRegister = []() {
|
|
OnnxExtraManager::get()->insert("RandomUniform",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxRandomUniformTransform));
|
|
OnnxExtraManager::get()->insert("RandomUniformLike",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxRandomUniformTransform));
|
|
OnnxExtraManager::get()->insert("RandomNormal",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxRandomUniformTransform));
|
|
OnnxExtraManager::get()->insert("RandomNormalLike",
|
|
std::shared_ptr<OnnxExtraManager::Transform>(new OnnxRandomUniformTransform));
|
|
return true;
|
|
}();
|
|
|
|
} // namespace Express
|
|
} // namespace MNN
|