From bc500172b04fb24d986748ad89f78fa1ba72c066 Mon Sep 17 00:00:00 2001 From: Håvard Pettersen Date: Mon, 17 Aug 2020 15:05:52 +0000 Subject: onnx wrapper --- eval/CMakeLists.txt | 1 + eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt | 9 + .../tensor/onnx_wrapper/onnx_wrapper_test.cpp | 119 +++++++++++ eval/src/vespa/eval/CMakeLists.txt | 1 + eval/src/vespa/eval/tensor/dense/CMakeLists.txt | 1 + eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp | 233 +++++++++++++++++++++ eval/src/vespa/eval/tensor/dense/onnx_wrapper.h | 84 ++++++++ 7 files changed, 448 insertions(+) create mode 100644 eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt create mode 100644 eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp create mode 100644 eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp create mode 100644 eval/src/vespa/eval/tensor/dense/onnx_wrapper.h diff --git a/eval/CMakeLists.txt b/eval/CMakeLists.txt index 3e81521550a..fe9d9985c6a 100644 --- a/eval/CMakeLists.txt +++ b/eval/CMakeLists.txt @@ -52,6 +52,7 @@ vespa_define_module( src/tests/tensor/direct_dense_tensor_builder src/tests/tensor/direct_sparse_tensor_builder src/tests/tensor/index_lookup_table + src/tests/tensor/onnx_wrapper src/tests/tensor/tensor_add_operation src/tests/tensor/tensor_address src/tests/tensor/tensor_conformance diff --git a/eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt b/eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt new file mode 100644 index 00000000000..9c92f75476c --- /dev/null +++ b/eval/src/tests/tensor/onnx_wrapper/CMakeLists.txt @@ -0,0 +1,9 @@ +# Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +vespa_add_executable(eval_onnx_wrapper_test_app TEST + SOURCES + onnx_wrapper_test.cpp + DEPENDS + vespaeval + GTest::GTest +) +vespa_add_test(NAME eval_onnx_wrapper_test_app COMMAND eval_onnx_wrapper_test_app) diff --git a/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp new file mode 100644 index 00000000000..07e065f9e39 --- /dev/null +++ b/eval/src/tests/tensor/onnx_wrapper/onnx_wrapper_test.cpp @@ -0,0 +1,119 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include +#include +#include +#include +#include + +using namespace vespalib::eval; +using namespace vespalib::tensor; + +using vespalib::make_string_short::fmt; + +std::string get_source_dir() { + const char *dir = getenv("SOURCE_DIRECTORY"); + return (dir ? dir : "."); +} +std::string source_dir = get_source_dir(); +std::string vespa_dir = source_dir + "/" + "../../../../.."; +std::string simple_model = vespa_dir + "/" + "model-integration/src/test/models/onnx/simple/simple.onnx"; + +vespalib::string to_str(const std::vector &dim_sizes) { + vespalib::string res; + for (size_t dim_size: dim_sizes) { + if (dim_size == 0) { + res += "[]"; + } else { + res += fmt("[%zu]", dim_size); + } + } + return res; +} + +vespalib::string to_str(OnnxWrapper::TensorInfo::ElementType element_type) { + if (element_type == OnnxWrapper::TensorInfo::ElementType::FLOAT) { + return "float"; + } + if (element_type == OnnxWrapper::TensorInfo::ElementType::DOUBLE) { + return "double"; + } + return "???"; +} + +void dump_info(const char *ctx, const std::vector &info) { + fprintf(stderr, "%s:\n", ctx); + for (size_t i = 0; i < info.size(); ++i) { + fprintf(stderr, " %s[%zu]: '%s' %s%s\n", ctx, i, info[i].name.c_str(), + to_str(info[i].elements).c_str(),to_str(info[i].dimensions).c_str()); + } +} + +TEST(OnnxWrapperTest, onnx_model_can_be_inspected) +{ + OnnxWrapper wrapper(simple_model, OnnxWrapper::Optimize::DISABLE); + dump_info("inputs", wrapper.inputs()); + dump_info("outputs", wrapper.outputs()); + ASSERT_EQ(wrapper.inputs().size(), 3); + ASSERT_EQ(wrapper.outputs().size(), 1); + //------------------------------------------------------------------------- + EXPECT_EQ( wrapper.inputs()[0].name, "query_tensor"); + EXPECT_EQ(to_str(wrapper.inputs()[0].dimensions), "[1][4]"); + EXPECT_EQ(to_str(wrapper.inputs()[0].elements), "float"); + //------------------------------------------------------------------------- + EXPECT_EQ( wrapper.inputs()[1].name, "attribute_tensor"); + EXPECT_EQ(to_str(wrapper.inputs()[1].dimensions), "[4][1]"); + EXPECT_EQ(to_str(wrapper.inputs()[1].elements), "float"); + //------------------------------------------------------------------------- + EXPECT_EQ( wrapper.inputs()[2].name, "bias_tensor"); + EXPECT_EQ(to_str(wrapper.inputs()[2].dimensions), "[1][1]"); + EXPECT_EQ(to_str(wrapper.inputs()[2].elements), "float"); + //------------------------------------------------------------------------- + EXPECT_EQ( wrapper.outputs()[0].name, "output"); + EXPECT_EQ(to_str(wrapper.outputs()[0].dimensions), "[1][1]"); + EXPECT_EQ(to_str(wrapper.outputs()[0].elements), "float"); +} + +TEST(OnnxWrapperTest, onnx_model_can_be_evaluated) +{ + OnnxWrapper wrapper(simple_model, OnnxWrapper::Optimize::ENABLE); + + ValueType query_type = ValueType::from_spec("tensor(a[1],b[4])"); + std::vector query_values({1.0, 2.0, 3.0, 4.0}); + DenseTensorView query(query_type, TypedCells(query_values)); + EXPECT_TRUE(wrapper.inputs()[0].is_compatible(query_type)); + EXPECT_FALSE(wrapper.inputs()[1].is_compatible(query_type)); + EXPECT_FALSE(wrapper.inputs()[2].is_compatible(query_type)); + + ValueType attribute_type = ValueType::from_spec("tensor(a[4],b[1])"); + std::vector attribute_values({5.0, 6.0, 7.0, 8.0}); + DenseTensorView attribute(attribute_type, TypedCells(attribute_values)); + EXPECT_FALSE(wrapper.inputs()[0].is_compatible(attribute_type)); + EXPECT_TRUE(wrapper.inputs()[1].is_compatible(attribute_type)); + EXPECT_FALSE(wrapper.inputs()[2].is_compatible(attribute_type)); + + ValueType bias_type = ValueType::from_spec("tensor(a[1],b[1])"); + std::vector bias_values({9.0}); + DenseTensorView bias(bias_type, TypedCells(bias_values)); + EXPECT_FALSE(wrapper.inputs()[0].is_compatible(bias_type)); + EXPECT_FALSE(wrapper.inputs()[1].is_compatible(bias_type)); + EXPECT_TRUE(wrapper.inputs()[2].is_compatible(bias_type)); + + MutableDenseTensorView output(wrapper.outputs()[0].make_compatible_type()); + EXPECT_EQ(output.fast_type().to_spec(), "tensor(d0[1],d1[1])"); + + OnnxWrapper::Params params; + params.bind(0, query); + params.bind(1, attribute); + params.bind(2, bias); + auto result = wrapper.eval(params); + + EXPECT_EQ(result.num_values(), 1); + result.get(0, output); + auto cells = output.cellsRef(); + EXPECT_EQ(cells.type, ValueType::CellType::FLOAT); + EXPECT_EQ(cells.size, 1); + EXPECT_EQ(cells.get(0), 79.0); +} + +GTEST_MAIN_RUN_ALL_TESTS() diff --git a/eval/src/vespa/eval/CMakeLists.txt b/eval/src/vespa/eval/CMakeLists.txt index c28643e605e..04f151f7ced 100644 --- a/eval/src/vespa/eval/CMakeLists.txt +++ b/eval/src/vespa/eval/CMakeLists.txt @@ -12,6 +12,7 @@ vespa_add_library(vespaeval $ INSTALL lib64 DEPENDS + onnxruntime ${VESPA_LLVM_LIB} ) diff --git a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt index c4b8138148c..b4e849a1dde 100644 --- a/eval/src/vespa/eval/tensor/dense/CMakeLists.txt +++ b/eval/src/vespa/eval/tensor/dense/CMakeLists.txt @@ -30,6 +30,7 @@ vespa_add_library(eval_tensor_dense OBJECT dense_xw_product_function.cpp index_lookup_table.cpp mutable_dense_tensor_view.cpp + onnx_wrapper.cpp typed_cells.cpp typed_dense_tensor_builder.cpp vector_from_doubles_function.cpp diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp new file mode 100644 index 00000000000..fa0379473c9 --- /dev/null +++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.cpp @@ -0,0 +1,233 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "onnx_wrapper.h" +#include +#include "dense_tensor_view.h" +#include "mutable_dense_tensor_view.h" +#include +#include +#include +#include +#include +#include + +using vespalib::eval::ValueType; +using vespalib::make_string_short::fmt; + +namespace vespalib::tensor { + +namespace { + +ValueType::CellType as_cell_type(OnnxWrapper::TensorInfo::ElementType type) { + if (type == OnnxWrapper::TensorInfo::ElementType::FLOAT) { + return ValueType::CellType::FLOAT; + } + if (type == OnnxWrapper::TensorInfo::ElementType::DOUBLE) { + return ValueType::CellType::DOUBLE; + } + abort(); +} + +auto convert_optimize(OnnxWrapper::Optimize optimize) { + if (optimize == OnnxWrapper::Optimize::ENABLE) { + return ORT_ENABLE_ALL; + } else { + assert(optimize == OnnxWrapper::Optimize::DISABLE); + return ORT_DISABLE_ALL; + } +} + +class OnnxString { +private: + static Ort::AllocatorWithDefaultOptions _alloc; + char *_str; + void cleanup() { + if (_str != nullptr) { + _alloc.Free(_str); + _str = nullptr; + } + } + OnnxString(char *str) : _str(str) {} +public: + OnnxString(const OnnxString &rhs) = delete; + OnnxString &operator=(const OnnxString &rhs) = delete; + OnnxString(OnnxString &&rhs) : _str(rhs._str) { + rhs._str = nullptr; + } + OnnxString &operator=(OnnxString &&rhs) { + cleanup(); + _str = rhs._str; + rhs._str = nullptr; + return *this; + } + const char *get() const { return _str; } + ~OnnxString() { cleanup(); } + static OnnxString get_input_name(const Ort::Session &session, size_t idx) { + return OnnxString(session.GetInputName(idx, _alloc)); + } + static OnnxString get_output_name(const Ort::Session &session, size_t idx) { + return OnnxString(session.GetOutputName(idx, _alloc)); + } +}; +Ort::AllocatorWithDefaultOptions OnnxString::_alloc; + +std::vector make_dimensions(const std::vector &shape) { + std::vector result; + for (int64_t size: shape) { + result.push_back(std::max(size, 0L)); + } + return result; +} + +OnnxWrapper::TensorInfo::ElementType make_element_type(ONNXTensorElementDataType element_type) { + if (element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT) { + return OnnxWrapper::TensorInfo::ElementType::FLOAT; + } else if (element_type == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE) { + return OnnxWrapper::TensorInfo::ElementType::DOUBLE; + } else { + return OnnxWrapper::TensorInfo::ElementType::UNKNOWN; + } +} + +OnnxWrapper::TensorInfo make_tensor_info(const OnnxString &name, const Ort::TypeInfo &type_info) { + auto tensor_info = type_info.GetTensorTypeAndShapeInfo(); + auto shape = tensor_info.GetShape(); + auto element_type = tensor_info.GetElementType(); + return OnnxWrapper::TensorInfo{vespalib::string(name.get()), make_dimensions(shape), make_element_type(element_type)}; +} + +} + +bool +OnnxWrapper::TensorInfo::is_compatible(const eval::ValueType &type) const +{ + if ((elements == ElementType::UNKNOWN) || dimensions.empty()) { + return false; + } + if (type.cell_type() != as_cell_type(elements)) { + return false; + } + if (type.dimensions().size() != dimensions.size()) { + return false; + } + for (size_t i = 0; i < dimensions.size(); ++i) { + if (type.dimensions()[i].size != dimensions[i]) { + return false; + } + } + return true; +} + +eval::ValueType +OnnxWrapper::TensorInfo::make_compatible_type() const +{ + if ((elements == ElementType::UNKNOWN) || dimensions.empty()) { + return ValueType::error_type(); + } + std::vector dim_list; + for (size_t dim_size: dimensions) { + if ((dim_size == 0) || (dim_list.size() > 9)) { + return ValueType::error_type(); + } + dim_list.emplace_back(fmt("d%zu", dim_list.size()), dim_size); + } + return ValueType::tensor_type(std::move(dim_list), as_cell_type(elements)); +} + +OnnxWrapper::TensorInfo::~TensorInfo() = default; + +OnnxWrapper::Shared::Shared() + : _env(ORT_LOGGING_LEVEL_WARNING, "vespa-onnx-wrapper") +{ +} + +void +OnnxWrapper::Params::bind(size_t idx, const DenseTensorView &src) +{ + assert(idx == values.size()); + std::vector dim_sizes; + for (const auto &dim: src.fast_type().dimensions()) { + dim_sizes.push_back(dim.size); + } + auto memory_info = Ort::MemoryInfo::CreateCpu(OrtArenaAllocator, OrtMemTypeDefault); + if (src.fast_type().cell_type() == ValueType::CellType::FLOAT) { + // NB: create requires non-const input + auto cells = unconstify(src.cellsRef().typify()); + values.push_back(Ort::Value::CreateTensor(memory_info, cells.begin(), cells.size(), dim_sizes.data(), dim_sizes.size())); + } else if (src.fast_type().cell_type() == ValueType::CellType::DOUBLE) { + // NB: create requires non-const input + auto cells = unconstify(src.cellsRef().typify()); + values.push_back(Ort::Value::CreateTensor(memory_info, cells.begin(), cells.size(), dim_sizes.data(), dim_sizes.size())); + } +} + +void +OnnxWrapper::Result::get(size_t idx, MutableDenseTensorView &dst) +{ + assert(values[idx].IsTensor()); + auto meta = values[idx].GetTensorTypeAndShapeInfo(); + if (dst.fast_type().cell_type() == ValueType::CellType::FLOAT) { + assert(meta.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_FLOAT); + ConstArrayRef cells(values[idx].GetTensorMutableData(), meta.GetElementCount()); + dst.setCells(TypedCells(cells)); + } else if (dst.fast_type().cell_type() == ValueType::CellType::DOUBLE) { + assert(meta.GetElementType() == ONNX_TENSOR_ELEMENT_DATA_TYPE_DOUBLE); + ConstArrayRef cells(values[idx].GetTensorMutableData(), meta.GetElementCount()); + dst.setCells(TypedCells(cells)); + } +} + +OnnxWrapper::Shared & +OnnxWrapper::Shared::get() { + static Shared shared; + return shared; +} + +void +OnnxWrapper::extract_meta_data() +{ + Ort::AllocatorWithDefaultOptions allocator; + size_t num_inputs = _session.GetInputCount(); + for (size_t i = 0; i < num_inputs; ++i) { + _inputs.push_back(make_tensor_info(OnnxString::get_input_name(_session, i), _session.GetInputTypeInfo(i))); + } + size_t num_outputs = _session.GetOutputCount(); + for (size_t i = 0; i < num_outputs; ++i) { + _outputs.push_back(make_tensor_info(OnnxString::get_output_name(_session, i), _session.GetOutputTypeInfo(i))); + } + for (const auto &input: _inputs) { + _input_name_refs.push_back(input.name.c_str()); + } + for (const auto &output: _outputs) { + _output_name_refs.push_back(output.name.c_str()); + } +} + +OnnxWrapper::OnnxWrapper(const vespalib::string &model_file, Optimize optimize) + : _shared(Shared::get()), + _options(), + _session(nullptr), + _inputs(), + _outputs(), + _input_name_refs(), + _output_name_refs() +{ + _options.SetIntraOpNumThreads(1); + _options.SetInterOpNumThreads(1); + _options.SetGraphOptimizationLevel(convert_optimize(optimize)); + _session = Ort::Session(_shared.env(), model_file.c_str(), _options); + extract_meta_data(); +} + +OnnxWrapper::~OnnxWrapper() = default; + +OnnxWrapper::Result +OnnxWrapper::eval(const Params ¶ms) +{ + assert(params.values.size() == _inputs.size()); + Ort::RunOptions run_opts(nullptr); + return Result(_session.Run(run_opts, _input_name_refs.data(), params.values.data(), _inputs.size(), + _output_name_refs.data(), _outputs.size())); +} + +} diff --git a/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h new file mode 100644 index 00000000000..67a64f2d318 --- /dev/null +++ b/eval/src/vespa/eval/tensor/dense/onnx_wrapper.h @@ -0,0 +1,84 @@ +// Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +#include +#include +#include +#include + +namespace vespalib::tensor { + +class DenseTensorView; +class MutableDenseTensorView; + +/** + * Wrapper around an ONNX model handeled by onnxruntime. + **/ +class OnnxWrapper { +public: + // model optimization + enum class Optimize { ENABLE, DISABLE }; + + // information about a single input or output tensor + struct TensorInfo { + enum class ElementType { FLOAT, DOUBLE, UNKNOWN }; + vespalib::string name; + std::vector dimensions; + ElementType elements; + bool is_compatible(const eval::ValueType &type) const; + eval::ValueType make_compatible_type() const; + ~TensorInfo(); + }; + + // used to build model parameters + class Params { + friend class OnnxWrapper; + private: + std::vector values; + public: + Params() : values() {} + void bind(size_t idx, const DenseTensorView &src); + }; + + // used to inspect model results + class Result { + friend class OnnxWrapper; + private: + std::vector values; + Result(std::vector values_in) : values(std::move(values_in)) {} + public: + size_t num_values() const { return values.size(); } + void get(size_t idx, MutableDenseTensorView &dst); + }; + +private: + // common stuff shared between model sessions + class Shared { + private: + Ort::Env _env; + Shared(); + public: + static Shared &get(); + Ort::Env &env() { return _env; } + }; + + Shared &_shared; + Ort::SessionOptions _options; + Ort::Session _session; + std::vector _inputs; + std::vector _outputs; + std::vector _input_name_refs; + std::vector _output_name_refs; + + void extract_meta_data(); + +public: + OnnxWrapper(const vespalib::string &model_file, Optimize optimize); + ~OnnxWrapper(); + const std::vector &inputs() const { return _inputs; } + const std::vector &outputs() const { return _outputs; } + Result eval(const Params ¶ms); // NB: Run requires non-const session +}; + +} -- cgit v1.2.3