From 3f204d191b0e4fb71873e71d83cbeb3452a9e1ce Mon Sep 17 00:00:00 2001 From: Rachel Guo <35738743+YUNQIUGUO@users.noreply.github.com> Date: Tue, 18 May 2021 15:12:41 -0700 Subject: [PATCH] [CoreML EP] Add Squeeze Op support (#7730) * add squeeze op builder initial * fix mistakes * modify * enable UT passed and minor refine * minor formatting * address comment Co-authored-by: rachguo --- .../builders/impl/squeeze_op_builder.cc | 104 ++++++++++++++++++ .../coreml/builders/op_builder_factory.cc | 4 + .../coreml/builders/op_builder_factory.h | 1 + .../providers/cpu/tensor/squeeze_op_test.cc | 46 +++++--- 4 files changed, 137 insertions(+), 18 deletions(-) create mode 100644 onnxruntime/core/providers/coreml/builders/impl/squeeze_op_builder.cc diff --git a/onnxruntime/core/providers/coreml/builders/impl/squeeze_op_builder.cc b/onnxruntime/core/providers/coreml/builders/impl/squeeze_op_builder.cc new file mode 100644 index 0000000000..bb9d61c11b --- /dev/null +++ b/onnxruntime/core/providers/coreml/builders/impl/squeeze_op_builder.cc @@ -0,0 +1,104 @@ +// Copyright (c) Microsoft Corporation. All rights reserved. +// Licensed under the MIT License. +#include + +#include "core/providers/common.h" +#include "core/providers/shared/utils/utils.h" +#include "core/providers/coreml/builders/model_builder.h" +#include "core/providers/coreml/builders/op_builder_factory.h" + +#include "base_op_builder.h" + +namespace onnxruntime { +namespace coreml { + +class SqueezeOpBuilder : public BaseOpBuilder { + // Add operator related + public: + void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override; + + private: + Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node, + const logging::Logger& logger) const override ORT_MUST_USE_RESULT; + + // Operator support related + private: + bool IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node, + const logging::Logger& logger) const override; +}; + +// Add operator related + +void SqueezeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const { + if (node.SinceVersion() > 12 && node.InputDefs().size() > 1) { + model_builder.AddInitializerToSkip(node.InputDefs()[1]->Name()); + } +} + +/* static */ std::vector GetAxes(ModelBuilder& model_builder, const Node& node) { + std::vector axes; + // Squeeze opset 13 use input as axes + if (node.SinceVersion() > 12) { + // If axes is not provided, return an empty axes as default to squeeze all + if (node.InputDefs().size() > 1) { + const auto& initializers(model_builder.GetInitializerTensors()); + const auto& axes_tensor = *initializers.at(node.InputDefs()[1]->Name()); + const int64_t* raw_axes = GetTensorInt64Data(axes_tensor); + const auto size = SafeInt(axes_tensor.dims()[0]); + axes.resize(size); + for (size_t i = 0; i < size; i++) { + axes[i] = raw_axes[i]; + } + } + } else { + NodeAttrHelper helper(node); + axes = helper.Get("axes", std::vector()); + } + + return axes; +} + +Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, + const Node& node, + const logging::Logger& /* logger */) const { + std::unique_ptr layer = CreateNNLayer(node); + + auto* coreml_squeeze = layer->mutable_squeeze(); + std::vector axes = GetAxes(model_builder, node); + if (axes.empty()) { + coreml_squeeze->set_squeezeall(true); + } else { + *coreml_squeeze->mutable_axes() = {axes.cbegin(), axes.cend()}; + coreml_squeeze->set_squeezeall(false); + } + + *layer->mutable_input()->Add() = node.InputDefs()[0]->Name(); + *layer->mutable_output()->Add() = node.OutputDefs()[0]->Name(); + + model_builder.AddLayer(std::move(layer)); + return Status::OK(); +} + +// Operator support related + +bool SqueezeOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& initializers, const Node& node, + const logging::Logger& /*logger*/) const { + // Squeeze opset 13 uses input 1 as axes, if we have input 1 then it needs to be an initializer + if (node.SinceVersion() > 12 && node.InputDefs().size() > 1) { + const auto& axes_name = node.InputDefs()[1]->Name(); + if (!Contains(initializers, axes_name)) { + LOGS_DEFAULT(VERBOSE) << "Input axes of Squeeze must be known"; + return false; + } + } + + return true; +} + +void CreateSqueezeOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations) { + op_registrations.builders.push_back(std::make_unique()); + op_registrations.op_builder_map.emplace(op_type, op_registrations.builders.back().get()); +} + +} // namespace coreml +} // namespace onnxruntime diff --git a/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc b/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc index 4c15040b84..171f0b198a 100644 --- a/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc +++ b/onnxruntime/core/providers/coreml/builders/op_builder_factory.cc @@ -65,6 +65,10 @@ static OpBuilderRegistrations CreateOpBuilderRegistrations() { CreateClipOpBuilder("Clip", op_registrations); } + { // Squeeze + CreateSqueezeOpBuilder("Squeeze", op_registrations); + } + return op_registrations; } diff --git a/onnxruntime/core/providers/coreml/builders/op_builder_factory.h b/onnxruntime/core/providers/coreml/builders/op_builder_factory.h index f32f551dcc..2617c45861 100644 --- a/onnxruntime/core/providers/coreml/builders/op_builder_factory.h +++ b/onnxruntime/core/providers/coreml/builders/op_builder_factory.h @@ -30,5 +30,6 @@ void CreateClipOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_ void CreateActivationOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreatePoolOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); void CreateGemmOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); +void CreateSqueezeOpBuilder(const std::string& op_type, OpBuilderRegistrations& op_registrations); } // namespace coreml } // namespace onnxruntime \ No newline at end of file diff --git a/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc b/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc index 39561a6ba5..bb31ef0512 100644 --- a/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc +++ b/onnxruntime/test/providers/cpu/tensor/squeeze_op_test.cc @@ -117,15 +117,20 @@ TEST(SqueezeOpTest, SqueezeNegAxis_2) { } TEST(SqueezeOpTest, Squeeze_2_axes_input) { - OpTester test("Squeeze", 13); - test.AddInput("data", {1, 4, 1, 1, 2}, - std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); - test.AddInput("axes", {3}, std::vector{0, 2, 3}); - test.AddOutput("squeezed", {4, 2}, - std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); - // Incorrect precision for OpenVINO EP. Will be re-enabled after it's fixed - // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider}); + auto run_test = [](bool axes_is_initializer) { + OpTester test("Squeeze", 13); + test.AddInput("data", {1, 4, 1, 1, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + test.AddInput("axes", {3}, std::vector{0, 2, 3}, axes_is_initializer); + test.AddOutput("squeezed", {4, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + // Incorrect precision for OpenVINO EP. Will be re-enabled after it's fixed + // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider}); + }; + + run_test(false); + run_test(true); // COREML EP will need axes as an initializer } TEST(SqueezeOpTest, Squeeze_Empty_Axes_opset13) { @@ -137,17 +142,22 @@ TEST(SqueezeOpTest, Squeeze_Empty_Axes_opset13) { } TEST(SqueezeOpTest, SqueezeNegAxis_axes_input) { - OpTester test("Squeeze", 13); - test.AddInput("data", {1, 4, 1, 1, 2}, - std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + auto run_test = [](bool axes_is_initializer) { + OpTester test("Squeeze", 13); + test.AddInput("data", {1, 4, 1, 1, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); - test.AddInput("axes", {3}, std::vector{0, -3, -2}); - test.AddOutput("squeezed", {4, 2}, - std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); + test.AddInput("axes", {3}, std::vector{0, -3, -2}, axes_is_initializer); + test.AddOutput("squeezed", {4, 2}, + std::vector{1.0f, 2.0f, 3.0f, 4.0f, 5.0f, 6.0f, 7.0f, 8.0f}); - // OpenVINO EP Incorrect precision. Will be re-enabled after its fixed - // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after - test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider}); + // OpenVINO EP Incorrect precision. Will be re-enabled after its fixed + // TensorRT and OpenVINO dont support "axes" input in opset 13, re-enable after + test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kOpenVINOExecutionProvider, kTensorrtExecutionProvider}); + }; + + run_test(false); + run_test(true); // COREML EP will need axes as an initializer } // Add 4d input shape test, since NNAPI supports up to 4d input shape