diff --git a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc index f992316cd3..555c76c0d9 100644 --- a/onnxruntime/core/providers/cpu/cpu_execution_provider.cc +++ b/onnxruntime/core/providers/cpu/cpu_execution_provider.cc @@ -196,6 +196,12 @@ class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, Eye class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, float, IsNaN); class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, MLFloat16, IsNaN); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, Erf); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int64_t_int64_t_int64_t, OneHot); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, float_int64_t_int64_t, OneHot); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int64_t_string_int64_t, OneHot); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, float_string_int64_t, OneHot); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, float_float_float, OneHot); +class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, int64_t_int32_t_float, OneHot); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, MaxUnpool); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, Sinh); class ONNX_OPERATOR_KERNEL_CLASS_NAME(kCpuExecutionProvider, kOnnxDomain, 9, Cosh); @@ -387,6 +393,12 @@ void RegisterOnnxOperatorKernels(std::function fn) { fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); + fn(BuildKernel()); + fn(BuildKernel()); + fn(BuildKernel()); + fn(BuildKernel()); + fn(BuildKernel()); + fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); fn(BuildKernel()); diff --git a/onnxruntime/core/providers/cpu/tensor/onehot.cc b/onnxruntime/core/providers/cpu/tensor/onehot.cc new file mode 100644 index 0000000000..806b0c8485 --- /dev/null +++ b/onnxruntime/core/providers/cpu/tensor/onehot.cc @@ -0,0 +1,169 @@ +/* Copyright 2016 The TensorFlow Authors. All Rights Reserved. + +Licensed under the Apache License, Version 2.0 (the "License"); +you may not use this file except in compliance with the License. +You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + +Unless required by applicable law or agreed to in writing, software +distributed under the License is distributed on an "AS IS" BASIS, +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +See the License for the specific language governing permissions and +limitations under the License. +==============================================================================*/ +/* Modifications Copyright (c) Microsoft. */ + +#include "core/providers/cpu/tensor/onehot.h" + +// build\windows\debug\external\eigen3\unsupported\eigen\cxx11\src/Tensor/Tensor.h(76): +// warning C4554: '&': check operator precedence for possible error; use parentheses to clarify precedence +#ifdef _MSC_VER +#pragma warning(push) +#pragma warning(disable : 4554) +#endif +#include "core/util/eigen_common_wrapper.h" +#ifdef _MSC_VER +#pragma warning(pop) +#endif + +#include "core/platform/env.h" + +#define EIGEN_USE_THREADS + +using namespace ::onnxruntime::common; +using namespace std; + +namespace onnxruntime { +// spec: https://github.com/onnx/onnx/blob/master/docs/Operators.md#OneHot + +// T1: indices, T2: depth, T3: values +#define REG_TYPED_ONE_HOT_OP(types_str, in_type, out_type, depth_type) \ + ONNX_CPU_OPERATOR_TYPED_KERNEL( \ + OneHot, \ + 9, \ + types_str, \ + KernelDefBuilder() \ + .TypeConstraint("T1", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T2", DataTypeImpl::GetTensorType()) \ + .TypeConstraint("T3", DataTypeImpl::GetTensorType()), \ + OneHotOp); + +#define REG_ONE_HOT_OP(in_type, out_type, depth_type) \ + REG_TYPED_ONE_HOT_OP(in_type##_##out_type##_##depth_type, in_type, out_type, depth_type) + +REG_ONE_HOT_OP(int64_t, int64_t, int64_t); +REG_ONE_HOT_OP(float, int64_t, int64_t); +REG_ONE_HOT_OP(int64_t, string, int64_t); +REG_ONE_HOT_OP(float, string, int64_t); +REG_ONE_HOT_OP(float, float, float); // added this to satisfy onnx model tests +REG_ONE_HOT_OP(int64_t, int32_t, float); // added this to satisfy onnx model tests + +Status ValidateInputs(const Tensor* depth, + const Tensor* values) { + // validation scenarios + // depth should be scalar and > 0 + if (!depth->Shape().IsScalar()) { + return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Invalid argument for depth; it's not a scalar."); + } + + // values Rank 1 tensor containing exactly two elements + if (!(values->Shape().NumDimensions() == 1 && values->Shape().Size() == 2)) { + return Status(ONNXRUNTIME, INVALID_ARGUMENT, + "Invalid argument for values; either it's rank is more than 1" + " or it has more than 2 elements"); + } + + return Status::OK(); +} + +// Helper to define Tensor types given that the scalar is of type T. +template +struct EigenTensorTypes { + using EigenTensorMap = Eigen::TensorMap, Eigen::Aligned>; + using ConstEigenTensorMap = Eigen::TensorMap, Eigen::Aligned>; + using Scalar = Eigen::TensorMap, Eigen::RowMajor, IndexType>, Eigen::Aligned>; + using ConstScalar = Eigen::TensorMap, Eigen::RowMajor, IndexType>, Eigen::Aligned>; + using ConstMatrix = Eigen::TensorMap, Eigen::Aligned>; +}; + +namespace generator { +template +class OneGenerator { + public: + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE + OneGenerator(const typename EigenTensorTypes::ConstMatrix& indices, + const typename EigenTensorTypes::ConstScalar& on_value, + const typename EigenTensorTypes::ConstScalar& off_value) + : indices_(indices), on_value_(on_value), off_value_(off_value) {} + + EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE out_type + operator()(const Eigen::array& pre_depth_suff) const { + return (indices_(pre_depth_suff[0], pre_depth_suff[2]) == pre_depth_suff[1]) + ? on_value_() + : off_value_(); + } + + private: + const typename EigenTensorTypes::ConstMatrix indices_; + const typename EigenTensorTypes::ConstScalar on_value_; + const typename EigenTensorTypes::ConstScalar off_value_; +}; +} // namespace generator + +template +Status OneHotOp::Compute(OpKernelContext* p_op_kernel_context) const { + const auto* indices = p_op_kernel_context->Input(0); + const auto* depth = p_op_kernel_context->Input(1); + const auto* values = p_op_kernel_context->Input(2); + + ORT_RETURN_IF_ERROR(ValidateInputs(depth, values)); + + // prepare output shape + const auto* depth_data = depth->Data(); + const int64_t depth_val = static_cast(*depth_data); // As per spec in case 'depth' is of non-integer type, it will be casted to int64 before use. + if (depth_val <= 0) { + return Status(ONNXRUNTIME, INVALID_ARGUMENT, "Depth is negative."); + } + + const auto& indices_shape = indices->Shape(); + const auto& indices_dims = indices_shape.GetDims(); + const auto indices_num_dims = indices_shape.NumDimensions(); + std::vector output_shape(indices_shape.GetDims()); + output_shape.insert(axis_ == -1 ? output_shape.end() : output_shape.begin() + axis_, + depth_val); + + // allocate output + const auto* values_data = values->Data(); + Tensor* output = p_op_kernel_context->Output(0, TensorShape(output_shape)); + + const int64_t axis = (axis_ == -1) ? indices_num_dims : axis_; + int64_t prefix_dim_size = 1; + for (int64_t i = 0; i < axis; ++i) { + prefix_dim_size *= indices_dims[i]; + } + const int64_t suffix_dim_size = indices_shape.Size() / prefix_dim_size; + + // Split indices into matrix of size prefix_dim_size x suffix_dim_size + Eigen::array indices_dims_e = {{prefix_dim_size, suffix_dim_size}}; + const auto* indices_data = indices->Data(); + typename EigenTensorTypes::ConstEigenTensorMap indices_tensor_e(indices_data, indices_dims_e); + + // Split output into 3-Tensor of size: + // prefix_dim_size x depth x suffix_dim_size. + Eigen::array output_dims_e = {{prefix_dim_size, depth_val, suffix_dim_size}}; + auto* output_data = output->MutableData(); + typename EigenTensorTypes::EigenTensorMap output_tensor_e(output_data, output_dims_e); + + typename EigenTensorTypes::ConstScalar on_value_e(values_data + 1); + typename EigenTensorTypes::ConstScalar off_value_e(values_data); + + generator::OneGenerator generator(indices_tensor_e, on_value_e, off_value_e); + + // TODO potential optimization opportunity + // TODO figure out the eigen threadpool stuff for use here + output_tensor_e = output_tensor_e.generate(generator); + + return Status::OK(); +} +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/cpu/tensor/onehot.h b/onnxruntime/core/providers/cpu/tensor/onehot.h new file mode 100644 index 0000000000..a05731b777 --- /dev/null +++ b/onnxruntime/core/providers/cpu/tensor/onehot.h @@ -0,0 +1,31 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#pragma once + +#include "core/common/common.h" +#include "core/framework/op_kernel.h" + +namespace onnxruntime { + +template +class OneHotOp final : public OpKernel { + public: + explicit OneHotOp(const OpKernelInfo& op_kernel_info) : OpKernel(op_kernel_info) { + int64_t tmp_axis; + if (op_kernel_info.GetAttr("axis", &tmp_axis).IsOK()) { + if (tmp_axis < -1) { // as per spec it can be -1 or more + ORT_THROW("Value of axis is < -1"); + } + axis_ = tmp_axis; + } + } + + Status Compute(OpKernelContext* p_op_kernel_context) const override; + + private: + ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(OneHotOp); + + int64_t axis_ = -1; +}; +} // namespace onnxruntime diff --git a/onnxruntime/test/onnx/main.cc b/onnxruntime/test/onnx/main.cc index 93a64012db..6436fce3b9 100644 --- a/onnxruntime/test/onnx/main.cc +++ b/onnxruntime/test/onnx/main.cc @@ -310,8 +310,6 @@ int real_main(int argc, char* argv[]) { {"prelu_broadcast", "disable reason"}, {"prelu_example", "disable reason"}, {"upsample_nearest", "opset 9 not supported yet"}, - {"onehot_with_axis", "opset 9 not supported yet"}, - {"onehot_without_axis", "opset 9 not supported yet"}, // also has bug in current test re: output type. Spandan to fix. {"asinh", "opset 9 not supported yet"}, {"acosh", "opset 9 not supported yet"}, {"atanh", "opset 9 not supported yet"}, diff --git a/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc new file mode 100644 index 0000000000..152cb3bcb9 --- /dev/null +++ b/onnxruntime/test/providers/cpu/tensor/onehot_op_test.cc @@ -0,0 +1,148 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. + +#include "gtest/gtest.h" +#include "test/providers/provider_test_utils.h" + +using namespace std; + +namespace onnxruntime { +namespace test { + +TEST(OneHotOpTest, DefaultAxis) { + OpTester test("OneHot", 9); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 10}, {0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, + 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_float_float_float /*indices, output, depth*/) { + OpTester test("OneHot", 9); + test.AddInput("indices", {2, 3}, {1., 9., 8., 2., 4., 6.}); + test.AddInput("depth", {1}, {10.}); + test.AddInput("values", {2}, {0., 1.}); + test.AddOutput("output", {2, 3, 10}, {0., 1., 0., 0., 0., 0., 0., 0., 0., 0., + 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., + 0., 0., 0., 0., 0., 0., 0., 0., 1., 0., + 0., 0., 1., 0., 0., 0., 0., 0., 0., 0., + 0., 0., 0., 0., 1., 0., 0., 0., 0., 0., + 0., 0., 0., 0., 0., 0., 1., 0., 0., 0.,}); + test.Run(); +} + +TEST(OneHotOpTest, DefaultAxis_int64_int32_float /*indices, output, depth*/) { + OpTester test("OneHot", 9); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10.}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 10}, {0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, + 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,}); + test.Run(); +} + +TEST(OneHotOpTest, Axis_0) { + OpTester test("OneHot", 9); + int64_t axis = 0; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {10, 2, 3}, { 0, 0, 0, 0, 0, 0, + 1, 0, 0, 0, 0, 0, + 0, 0, 0, 1, 0, 0, + 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, + 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 1, + 0, 0, 0, 0, 0, 0, + 0, 0, 1, 0, 0, 0, + 0, 1, 0, 0, 0, 0,}); + test.Run(); +} + +TEST(OneHotOpTest, Axis_1) { + OpTester test("OneHot", 9); + int64_t axis = 1; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 10, 3}, { 0, 0, 0, + 1, 0, 0, + 0, 0, 0, + 0, 0, 0, + 0, 0, 0, + 0, 0, 0, + 0, 0, 0, + 0, 0, 0, + 0, 0, 1, + 0, 1, 0, + 0, 0, 0, + 0, 0, 0, + 1, 0, 0, + 0, 0, 0, + 0, 1, 0, + 0, 0, 0, + 0, 0, 1, + 0, 0, 0, + 0, 0, 0, + 0, 0, 0,}); + test.Run(); +} + +TEST(OneHotOpTest, Axis_2) { + OpTester test("OneHot", 9); + int64_t axis = 2; + test.AddAttribute("axis", axis); + test.AddInput("indices", {2, 3}, {1, 9, 8, 2, 4, 6}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 10}, { 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, + 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,}); + test.Run(); +} + +TEST(OneHotOpTest, FloatInt64) { + OpTester test("OneHot", 9); + test.AddInput("indices", {2, 3}, {1.f, 9.f, 8.f, 2.f, 4.f, 6.f}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {0, 1}); + test.AddOutput("output", {2, 3, 10}, {0, 1, 0, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 0, 0, 0, 1, + 0, 0, 0, 0, 0, 0, 0, 0, 1, 0, + 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 1, 0, 0, 0, 0, 0, + 0, 0, 0, 0, 0, 0, 1, 0, 0, 0,}); + test.Run(); +} + +TEST(OneHotOpTest, FloatString) { + OpTester test("OneHot", 9); + test.AddInput("indices", {2, 3}, {1.f, 9.f, 8.f, 2.f, 4.f, 6.f}); + test.AddInput("depth", {1}, {10}); + test.AddInput("values", {2}, {"off", "on"}); + test.AddOutput("output", {2, 3, 10}, {"off", "on", "off", "off", "off", "off", "off", "off", "off", "off", + "off", "off", "off", "off", "off", "off", "off", "off", "off", "on", + "off", "off", "off", "off", "off", "off", "off", "off", "on", "off", + "off", "off", "on", "off", "off", "off", "off", "off", "off", "off", + "off", "off", "off", "off", "on", "off", "off", "off", "off", "off", + "off", "off", "off", "off", "off", "off", "on", "off", "off", "off",}); + test.Run(); +} +} +} // namespace onnxruntime diff --git a/onnxruntime/test/python/onnx_backend_test_series.py b/onnxruntime/test/python/onnx_backend_test_series.py index ea6ea135bc..7dc01e3083 100644 --- a/onnxruntime/test/python/onnx_backend_test_series.py +++ b/onnxruntime/test/python/onnx_backend_test_series.py @@ -34,8 +34,6 @@ backend_test.exclude(r'(test_acosh_cpu.*' '|test_eyelike_with_dtype_cpu.*' '|test_eyelike_without_dtype_cpu.*' '|test_gru_seq_length_cpu.*' -'|test_onehot_with_axis_cpu.*' -'|test_onehot_without_axis_cpu.*' '|test_scan_sum_cpu.*' '|test_scatter_with_axis_cpu.*' '|test_scatter_without_axis_cpu.*'