wehub-resource-sync
264 lines
9.3 KiB
C++
264 lines
9.3 KiB
C++
#include <cmath>
|
|
#include <cstdio>
|
|
#include <fstream>
|
|
#include <memory>
|
|
#include <string>
|
|
#include <vector>
|
|
#include <MNN/Interpreter.hpp>
|
|
#include "MNN_generated.h"
|
|
#include "cli.hpp"
|
|
#include "onnx.pb.h"
|
|
|
|
static void addTensorShape(onnx::ValueInfoProto* valueInfo, const std::string& name, const std::vector<int>& dims) {
|
|
valueInfo->set_name(name);
|
|
auto* tensorType = valueInfo->mutable_type()->mutable_tensor_type();
|
|
tensorType->set_elem_type(onnx::TensorProto_DataType_FLOAT);
|
|
auto* shape = tensorType->mutable_shape();
|
|
for (int dim : dims) {
|
|
shape->add_dim()->set_dim_value(dim);
|
|
}
|
|
}
|
|
|
|
static void addFloatInitializer(onnx::GraphProto* graph, const std::string& name, const std::vector<int64_t>& dims,
|
|
const std::vector<float>& values) {
|
|
auto* tensor = graph->add_initializer();
|
|
tensor->set_name(name);
|
|
tensor->set_data_type(onnx::TensorProto_DataType_FLOAT);
|
|
for (auto dim : dims) {
|
|
tensor->add_dims(dim);
|
|
}
|
|
for (auto value : values) {
|
|
tensor->add_float_data(value);
|
|
}
|
|
}
|
|
|
|
static void addInt64Initializer(onnx::GraphProto* graph, const std::string& name, const std::vector<int64_t>& dims,
|
|
const std::vector<int64_t>& values) {
|
|
auto* tensor = graph->add_initializer();
|
|
tensor->set_name(name);
|
|
tensor->set_data_type(onnx::TensorProto_DataType_INT64);
|
|
for (auto dim : dims) {
|
|
tensor->add_dims(dim);
|
|
}
|
|
for (auto value : values) {
|
|
tensor->add_int64_data(value);
|
|
}
|
|
}
|
|
|
|
static bool saveModel(const onnx::ModelProto& model, const std::string& fileName) {
|
|
std::ofstream output(fileName, std::ios::binary | std::ios::trunc);
|
|
return model.SerializeToOstream(&output);
|
|
}
|
|
|
|
static bool compareVector(const float* got, const std::vector<float>& expected, float tolerance = 1e-5f) {
|
|
for (size_t i = 0; i < expected.size(); ++i) {
|
|
if (std::fabs(got[i] - expected[i]) > tolerance) {
|
|
std::fprintf(stderr, "mismatch at %zu, expect=%f, got=%f\n", i, expected[i], got[i]);
|
|
return false;
|
|
}
|
|
}
|
|
return true;
|
|
}
|
|
|
|
static bool runMNNModel(const std::string& modelPath,
|
|
const std::vector<std::pair<std::string, std::vector<float>>>& inputs,
|
|
const std::vector<float>& expectedOutput) {
|
|
std::unique_ptr<MNN::Interpreter> net(MNN::Interpreter::createFromFile(modelPath.c_str()));
|
|
if (!net) {
|
|
return false;
|
|
}
|
|
MNN::ScheduleConfig config;
|
|
config.type = MNN_FORWARD_CPU;
|
|
auto session = net->createSession(config);
|
|
if (!session) {
|
|
return false;
|
|
}
|
|
for (const auto& item : inputs) {
|
|
auto* inputTensor = net->getSessionInput(session, item.first.c_str());
|
|
if (!inputTensor) {
|
|
return false;
|
|
}
|
|
MNN::Tensor hostTensor(inputTensor, inputTensor->getDimensionType());
|
|
::memcpy(hostTensor.host<float>(), item.second.data(), item.second.size() * sizeof(float));
|
|
inputTensor->copyFromHostTensor(&hostTensor);
|
|
}
|
|
if (net->runSession(session) != MNN::NO_ERROR) {
|
|
return false;
|
|
}
|
|
auto* outputTensor = net->getSessionOutput(session, nullptr);
|
|
if (!outputTensor) {
|
|
return false;
|
|
}
|
|
MNN::Tensor hostOutput(outputTensor, outputTensor->getDimensionType());
|
|
outputTensor->copyToHostTensor(&hostOutput);
|
|
return compareVector(hostOutput.host<float>(), expectedOutput);
|
|
}
|
|
|
|
static bool convertOnnx(const std::string& onnxModel, const std::string& mnnModel) {
|
|
modelConfig config;
|
|
config.model = modelConfig::ONNX;
|
|
config.modelFile = onnxModel;
|
|
config.MNNModel = mnnModel;
|
|
config.keepInputFormat = true;
|
|
return MNN::Cli::convertModel(config);
|
|
}
|
|
|
|
static onnx::ModelProto makeConcatEinsumModel() {
|
|
onnx::ModelProto model;
|
|
model.set_ir_version(8);
|
|
model.mutable_opset_import()->Add()->set_version(13);
|
|
auto* graph = model.mutable_graph();
|
|
graph->set_name("ConcatEinsum");
|
|
|
|
addTensorShape(graph->add_input(), "x", {2, 3});
|
|
addTensorShape(graph->add_input(), "y", {2, 3});
|
|
addTensorShape(graph->add_output(), "out", {2, 3});
|
|
|
|
addFloatInitializer(graph, "weight", {2}, {1.5f, -0.5f});
|
|
addInt64Initializer(graph, "axes", {1}, {0});
|
|
|
|
auto* unsqueezeX = graph->add_node();
|
|
unsqueezeX->set_op_type("Unsqueeze");
|
|
unsqueezeX->add_input("x");
|
|
unsqueezeX->add_input("axes");
|
|
unsqueezeX->add_output("x_unsqueezed");
|
|
|
|
auto* unsqueezeY = graph->add_node();
|
|
unsqueezeY->set_op_type("Unsqueeze");
|
|
unsqueezeY->add_input("y");
|
|
unsqueezeY->add_input("axes");
|
|
unsqueezeY->add_output("y_unsqueezed");
|
|
|
|
auto* concat = graph->add_node();
|
|
concat->set_op_type("Concat");
|
|
concat->add_input("x_unsqueezed");
|
|
concat->add_input("y_unsqueezed");
|
|
concat->add_output("stacked");
|
|
auto* axis = concat->add_attribute();
|
|
axis->set_name("axis");
|
|
axis->set_i(0);
|
|
|
|
auto* einsum = graph->add_node();
|
|
einsum->set_op_type("Einsum");
|
|
einsum->add_input("weight");
|
|
einsum->add_input("stacked");
|
|
einsum->add_output("out");
|
|
auto* equation = einsum->add_attribute();
|
|
equation->set_name("equation");
|
|
equation->set_s("i,i...->...");
|
|
return model;
|
|
}
|
|
|
|
static onnx::ModelProto makeReduceEinsumModel() {
|
|
onnx::ModelProto model;
|
|
model.set_ir_version(8);
|
|
model.mutable_opset_import()->Add()->set_version(13);
|
|
auto* graph = model.mutable_graph();
|
|
graph->set_name("ReduceEinsum");
|
|
|
|
addTensorShape(graph->add_input(), "stacked", {2, 2, 3});
|
|
addTensorShape(graph->add_output(), "out", {2, 3});
|
|
addFloatInitializer(graph, "weight", {2}, {1.5f, -0.5f});
|
|
|
|
auto* einsum = graph->add_node();
|
|
einsum->set_op_type("Einsum");
|
|
einsum->add_input("weight");
|
|
einsum->add_input("stacked");
|
|
einsum->add_output("out");
|
|
auto* equation = einsum->add_attribute();
|
|
equation->set_name("equation");
|
|
equation->set_s("i,i...->...");
|
|
return model;
|
|
}
|
|
|
|
static onnx::ModelProto makeConcat3EinsumModel() {
|
|
onnx::ModelProto model;
|
|
model.set_ir_version(8);
|
|
model.mutable_opset_import()->Add()->set_version(13);
|
|
auto* graph = model.mutable_graph();
|
|
graph->set_name("Concat3Einsum");
|
|
|
|
addTensorShape(graph->add_input(), "x", {2, 3});
|
|
addTensorShape(graph->add_input(), "y", {2, 3});
|
|
addTensorShape(graph->add_input(), "z", {2, 3});
|
|
addTensorShape(graph->add_output(), "out", {2, 3});
|
|
|
|
addFloatInitializer(graph, "weight", {3}, {1.0f, -2.0f, 0.5f});
|
|
addInt64Initializer(graph, "axes", {1}, {0});
|
|
|
|
auto* unsqueezeX = graph->add_node();
|
|
unsqueezeX->set_op_type("Unsqueeze");
|
|
unsqueezeX->add_input("x");
|
|
unsqueezeX->add_input("axes");
|
|
unsqueezeX->add_output("x_unsqueezed");
|
|
|
|
auto* unsqueezeY = graph->add_node();
|
|
unsqueezeY->set_op_type("Unsqueeze");
|
|
unsqueezeY->add_input("y");
|
|
unsqueezeY->add_input("axes");
|
|
unsqueezeY->add_output("y_unsqueezed");
|
|
|
|
auto* unsqueezeZ = graph->add_node();
|
|
unsqueezeZ->set_op_type("Unsqueeze");
|
|
unsqueezeZ->add_input("z");
|
|
unsqueezeZ->add_input("axes");
|
|
unsqueezeZ->add_output("z_unsqueezed");
|
|
|
|
auto* concat = graph->add_node();
|
|
concat->set_op_type("Concat");
|
|
concat->add_input("x_unsqueezed");
|
|
concat->add_input("y_unsqueezed");
|
|
concat->add_input("z_unsqueezed");
|
|
concat->add_output("stacked");
|
|
auto* axis = concat->add_attribute();
|
|
axis->set_name("axis");
|
|
axis->set_i(0);
|
|
|
|
auto* einsum = graph->add_node();
|
|
einsum->set_op_type("Einsum");
|
|
einsum->add_input("weight");
|
|
einsum->add_input("stacked");
|
|
einsum->add_output("out");
|
|
auto* equation = einsum->add_attribute();
|
|
equation->set_name("equation");
|
|
equation->set_s("i,i...->...");
|
|
return model;
|
|
}
|
|
|
|
int main() {
|
|
const std::string concatOnnx = "/tmp/mnn_concat_einsum.onnx";
|
|
const std::string concatMnn = "/tmp/mnn_concat_einsum.mnn";
|
|
const std::string reduceOnnx = "/tmp/mnn_reduce_einsum.onnx";
|
|
const std::string reduceMnn = "/tmp/mnn_reduce_einsum.mnn";
|
|
const std::string concat3Onnx = "/tmp/mnn_concat3_einsum.onnx";
|
|
const std::string concat3Mnn = "/tmp/mnn_concat3_einsum.mnn";
|
|
|
|
bool ok = saveModel(makeConcatEinsumModel(), concatOnnx);
|
|
ok = saveModel(makeReduceEinsumModel(), reduceOnnx) && ok;
|
|
ok = saveModel(makeConcat3EinsumModel(), concat3Onnx) && ok;
|
|
|
|
ok = convertOnnx(concatOnnx, concatMnn) && ok;
|
|
ok = convertOnnx(reduceOnnx, reduceMnn) && ok;
|
|
ok = convertOnnx(concat3Onnx, concat3Mnn) && ok;
|
|
|
|
const std::vector<float> x = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f};
|
|
const std::vector<float> y = {6.0f, 5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
|
|
const std::vector<float> expected = {-1.5f, 0.5f, 2.5f, 4.5f, 6.5f, 8.5f};
|
|
ok = runMNNModel(concatMnn, {{"x", x}, {"y", y}}, expected) && ok;
|
|
|
|
const std::vector<float> stacked = {1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 6.0f, 5.0f, 4.0f, 3.0f, 2.0f, 1.0f};
|
|
ok = runMNNModel(reduceMnn, {{"stacked", stacked}}, expected) && ok;
|
|
|
|
const std::vector<float> z = {0.0f, 1.0f, 0.0f, 1.0f, 0.0f, 1.0f};
|
|
const std::vector<float> expected3 = {-11.0f, -7.5f, -5.0f, -1.5f, 1.0f, 4.5f};
|
|
ok = runMNNModel(concat3Mnn, {{"x", x}, {"y", y}, {"z", z}}, expected3) && ok;
|
|
|
|
::remove(concatOnnx.c_str());
|
|
::remove(concatMnn.c_str());
|
|
::remove(reduceOnnx.c_str());
|
|
::remove(reduceMnn.c_str());
|
|
::remove(concat3Onnx.c_str());
|
|
::remove(concat3Mnn.c_str());
|
|
return ok ? 0 : 1;
|
|
}
|