From b07b647f66b1b02514aae41e0710a698b73fed9e Mon Sep 17 00:00:00 2001 From: Scott McKay Date: Thu, 8 Jun 2023 13:56:06 +1000 Subject: [PATCH] Fix some issues with NNAPI Softmax (#16095) ### Description Update NNAPI Softmax to coerce to 2D when opset is < 13. This prevents the layout change to NHWC from breaking the implementation, as well as making it work correctly when the ONNX node's axis != 1. Add check for opset 13+ that axis is inner-most dimension as we don't currently handle any other value correctly. Update tests to add model to check NHWC layout, as well as 4D tests. We didn't notice the issues with the NNAPI EP as it was only processing input shapes that were 2D or 4D (which was overly restrictive as well). ### Motivation and Context #15949 --- .../builders/impl/softmax_op_builder.cc | 115 +++++++++++++----- .../test/providers/cpu/math/softmax_test.cc | 92 ++++++++++---- .../test/testdata/ort_github_issue_15949.onnx | Bin 0 -> 163 bytes 3 files changed, 157 insertions(+), 50 deletions(-) create mode 100644 onnxruntime/test/testdata/ort_github_issue_15949.onnx diff --git a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/softmax_op_builder.cc b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/softmax_op_builder.cc index b6007628fb..1e420fec80 100644 --- a/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/softmax_op_builder.cc +++ b/onnxruntime/core/providers/nnapi/nnapi_builtin/builders/impl/softmax_op_builder.cc @@ -63,16 +63,11 @@ Status SoftMaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, cons const auto& operand_indices(model_builder.GetOperandIndices()); const auto& operand_types(model_builder.GetOperandTypes()); const auto android_feature_level = model_builder.GetEffectiveFeatureLevel(); + + const auto& input = node_unit.Inputs()[0].node_arg.Name(); + const auto& output = node_unit.Outputs()[0].node_arg.Name(); + NodeAttrHelper helper(node_unit); - - auto input = node_unit.Inputs()[0].node_arg.Name(); - - // TODO: Needs fix. - if (android_feature_level < ANEURALNETWORKS_FEATURE_LEVEL_3) { - ORT_ENFORCE(model_builder.UseNCHW(), - "For Android API Level < 29 input for softmax needs to be NCHW."); - } - int32_t axis = helper.Get("axis", 1); // Check if the quantization scale and ZP are correct @@ -90,20 +85,68 @@ Status SoftMaxOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, cons y_scale = 1.f / 256; } - const auto& output = node_unit.Outputs()[0].node_arg.Name(); - float beta = 1.f; InlinedVector input_indices; input_indices.push_back(operand_indices.at(input)); + const float beta = 1.f; ADD_SCALAR_OPERAND(model_builder, input_indices, beta); - if (android_feature_level > ANEURALNETWORKS_FEATURE_LEVEL_2) { - // you can only specify axis for android api level 29+ - ADD_SCALAR_OPERAND(model_builder, input_indices, axis); + auto input_shape = shaper[input]; + if (axis < 0) { + axis = static_cast(HandleNegativeAxis(axis, input_shape.size())); } - const OperandType output_operand_type(operand_types.at(input).type, shaper[output], y_scale, y_zero_point); - ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SOFTMAX, input_indices, - {output}, {output_operand_type})); + // if opset < 13 we may need to manually coerce into 2D. we can skip this IFF it's already 2D and axis == 1. + // otherwise we need the coercion to create an input shape that works with axis == 1, which will work with any + // NNAPI version. + // e.g. if 2D and axis is 0 we coerce to shape {1 , dim0 + dim1} + // if 4D and axis is 2 we coerce to shape {dim0 + dim1, dim2 + dim3} + if (node_unit.SinceVersion() < 13 && !(input_shape.size() == 2 && axis == 1)) { + // Reshape to 2D based on axis + const SafeInt safe1(1); + uint32_t dim0 = std::accumulate(input_shape.cbegin(), input_shape.cbegin() + axis, safe1, std::multiplies()); + uint32_t dim1 = std::accumulate(input_shape.cbegin() + axis, input_shape.cend(), safe1, std::multiplies()); + Shape input2d_shape{dim0, dim1}; + + std::string shape2d_name = model_builder.GetUniqueName(node_unit.Name() + input + "_2D_shape"); + std::string reshape2d_output_name = model_builder.GetUniqueName(node_unit.Name() + input + "_2D"); + + ORT_RETURN_IF_ERROR(op_builder_helpers::AddNnapiReshape(model_builder, input, shape2d_name, + {narrow(dim0), narrow(dim1)}, + reshape2d_output_name)); + + input_indices[0] = operand_indices.at(reshape2d_output_name); // replace input 1 with 2d output + + std::string softmax2d_output_name = model_builder.GetUniqueName("softmax_" + reshape2d_output_name); + const OperandType output_operand_type(operand_types.at(input).type, input2d_shape, y_scale, y_zero_point); + shaper.AddShape(softmax2d_output_name, input2d_shape); + + ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SOFTMAX, input_indices, + {softmax2d_output_name}, {output_operand_type})); + + // add Reshape back to original shape + const Shape& output_shape = shaper[output]; + + // convert from uint32_t to int32_t + std::vector reshape_output_shape; + reshape_output_shape.reserve(output_shape.size()); + std::transform(output_shape.begin(), output_shape.end(), std::back_inserter(reshape_output_shape), + [](uint32_t dim) { return narrow(dim); }); + + std::string reshape_output_name = model_builder.GetUniqueName(node_unit.Name() + output + "_shape"); + ORT_RETURN_IF_ERROR(op_builder_helpers::AddNnapiReshape(model_builder, softmax2d_output_name, + reshape_output_name, reshape_output_shape, + output)); + + } else { + if (android_feature_level > ANEURALNETWORKS_FEATURE_LEVEL_2) { + // you can only specify axis for android api level 29+ + ADD_SCALAR_OPERAND(model_builder, input_indices, axis); + } + + const OperandType output_operand_type(operand_types.at(input).type, shaper[output], y_scale, y_zero_point); + ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SOFTMAX, input_indices, + {output}, {output_operand_type})); + } return Status::OK(); } @@ -119,21 +162,35 @@ bool SoftMaxOpBuilder::IsOpSupportedImpl(const InitializedTensorSet& /* initiali if (!GetShape(node_unit.Inputs()[0].node_arg, input_shape)) return false; - const auto input_size = input_shape.size(); - if (input_size != 2 && input_size != 4) { - LOGS_DEFAULT(VERBOSE) << "SoftMax only support 2d/4d shape, input is " - << input_size << "d shape"; - return false; - } + if (node_unit.SinceVersion() < 13) { + // if opset < 13 the ONNX spec coerces to 2D based on axis, so we will manually do that when adding to the model + // and use axis of 1. + } else { + const auto input_size = narrow(input_shape.size()); + + if (input_size > 4) { + LOGS_DEFAULT(VERBOSE) << "Softmax only supports maximum 4d input. input is " << input_size << "d"; + return false; + } - if (params.android_feature_level < ANEURALNETWORKS_FEATURE_LEVEL_3) { NodeAttrHelper helper(node_unit); int32_t axis = helper.Get("axis", 1); - if (axis != 1) { - LOGS_DEFAULT(VERBOSE) - << "SoftMax only support axis 1 on Android API level: " << params.android_feature_level - << " input axis: " << axis; - return false; + if (axis < 0) { + axis = static_cast(HandleNegativeAxis(axis, input_shape.size())); + } + + if (params.android_feature_level < ANEURALNETWORKS_FEATURE_LEVEL_3) { + // 2D or 4D input is supported with axis of the last dim + if (input_size != 2 && input_size != 4) { + LOGS_DEFAULT(VERBOSE) << "Softmax only support 2d or 4d shape, input has " << input_size << "d shape"; + return false; + } + + if (axis != input_size - 1) { + LOGS_DEFAULT(VERBOSE) << "Softmax only supports axis of the last dim on Android API level " + << params.android_feature_level << ". input axis: " << axis; + return false; + } } } diff --git a/onnxruntime/test/providers/cpu/math/softmax_test.cc b/onnxruntime/test/providers/cpu/math/softmax_test.cc index 94479a64d5..e64abdd112 100644 --- a/onnxruntime/test/providers/cpu/math/softmax_test.cc +++ b/onnxruntime/test/providers/cpu/math/softmax_test.cc @@ -1,12 +1,14 @@ // Copyright (c) Microsoft Corporation. All rights reserved. // Licensed under the MIT License. +#include #include "gmock/gmock.h" #include "gtest/gtest.h" + +#include "core/session/environment.h" #include "test/providers/provider_test_utils.h" #include "test/common/cuda_op_test_utils.h" #include "test/common/dnnl_op_test_utils.h" -#include namespace onnxruntime { namespace test { @@ -123,9 +125,11 @@ TEST(SoftmaxOperator, LargeNumber) { } // np.random.seed(123) # Use a seed so we can replicate the input and expected values here and in python +// # create 60 input values // x = np.abs(np.random.randn(3, 4, 5).astype(np.float32)) -static std::vector three_dimensions = {3, 4, 5}; -static std::vector x_vals_3dims = { +static const std::vector three_dimensions = {3, 4, 5}; +static const std::vector four_dimensions = {1, 3, 4, 5}; +static const std::vector input_vals_60 = { 1.0856307f, 0.99734545f, 0.2829785f, 1.5062947f, 0.5786002f, 1.6514366f, 2.4266791f, 0.42891264f, 1.2659363f, 0.8667404f, 0.6788862f, 0.09470897f, 1.4913896f, 0.638902f, 0.44398195f, @@ -141,8 +145,8 @@ static std::vector x_vals_3dims = { 1.2940853f, 1.0387882f, 1.7437122f, 0.79806274f, 0.02968323f, 1.0693159f, 0.8907064f, 1.7548862f, 1.4956441f, 1.0693927f}; -TEST(SoftmaxOperator, ThreeDimsAxis0) { - // x = +TEST(SoftmaxOperator, ThreeAndFourDimsAxis0) { + // x = // node = onnx.helper.make_node('Softmax', inputs = ['x'], outputs = ['y'], axis = 0) // y = softmax_2d(x.reshape(1, 60)).reshape(3, 4, 5) // expect(node, inputs = [x], outputs = [y], @@ -164,11 +168,17 @@ TEST(SoftmaxOperator, ThreeDimsAxis0) { 0.017545262f, 0.0135920765f, 0.027506188f, 0.010684152f, 0.0049549243f, 0.01401341f, 0.011721271f, 0.027815264f, 0.021463264f, 0.014014485f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 0, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); // Axis=0 is not supported by TensorRT + RunTest(input_vals_60, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 0, + // axis=0 is not supported by TensorRT + {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); + + RunTest(input_vals_60, expected_vals, four_dimensions, /*opset*/ 7, /*axis*/ 0, + // axis=0 is not supported by TensorRT + {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); } -TEST(SoftmaxOperator, ThreeDimsAxis1) { - // x = +TEST(SoftmaxOperator, ThreeAndFourDimsSecondLastAxis) { + // x = // node = onnx.helper.make_node('Softmax', inputs = ['x'], outputs = ['y'], axis = 1) // y = softmax_2d(x.reshape(3, 20)).reshape(3, 4, 5) // expect(node, inputs = [x], outputs = [y], @@ -190,10 +200,14 @@ TEST(SoftmaxOperator, ThreeDimsAxis1) { 0.050680935f, 0.03926183f, 0.079453886f, 0.030862054f, 0.014312706f, 0.040478885f, 0.033857856f, 0.080346674f, 0.06199841f, 0.040481992f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 1, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); + RunTest(input_vals_60, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 1, + {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); + + RunTest(input_vals_60, expected_vals, four_dimensions, /*opset*/ 7, /*axis*/ 2, + {kTensorrtExecutionProvider, kOpenVINOExecutionProvider, kDnnlExecutionProvider}); } -TEST(SoftmaxOperator, ThreeDimsAxis1_opset13) { +TEST(SoftmaxOperator, ThreeAndFourDimsSecondLastAxis_opset13) { // For the same input, opset-13's behavior is different from an earlier opset // and we see different expected results for the same test input @@ -213,12 +227,15 @@ TEST(SoftmaxOperator, ThreeDimsAxis1_opset13) { 0.37181312f, 0.12944824f, 0.3946307f, 0.19975942f, 0.0699691f, 0.29696727f, 0.11163106f, 0.39906505f, 0.4012943f, 0.1979003f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 13, /*axis*/ 1, + RunTest(input_vals_60, expected_vals, three_dimensions, /*opset*/ 13, /*axis*/ 1, + {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO doesn't support opset-13 yet + + RunTest(input_vals_60, expected_vals, four_dimensions, /*opset*/ 13, /*axis*/ 2, {kTensorrtExecutionProvider, kOpenVINOExecutionProvider}); // OpenVINO doesn't support opset-13 yet } -TEST(SoftmaxOperator, ThreeDimsAxis2) { - // x = +TEST(SoftmaxOperator, ThreeAndFourDimsLastAxis) { + // x = // node = onnx.helper.make_node('Softmax', inputs = ['x'], outputs = ['y'], axis = 2) // y = softmax_2d(x.reshape(12, 5)).reshape(3, 4, 5) // expect(node, inputs = [x], outputs = [y], @@ -240,10 +257,11 @@ TEST(SoftmaxOperator, ThreeDimsAxis2) { 0.23619612f, 0.1829779f, 0.37029108f, 0.14383113f, 0.0667037f, 0.15740506f, 0.13165872f, 0.31243387f, 0.24108529f, 0.15741715f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 2); + RunTest(input_vals_60, expected_vals, three_dimensions, /*opset*/ 7, /*axis*/ 2); + RunTest(input_vals_60, expected_vals, four_dimensions, /*opset*/ 7, /*axis*/ 3); } -TEST(SoftmaxOperator, ThreeDimsAxis2_opset13) { +TEST(SoftmaxOperator, ThreeAndFourDimsLastAxis_opset13) { std::vector expected_vals = { 0.22277209f, 0.20394778f, 0.09983283f, 0.33927578f, 0.13417149f, 0.21729809f, 0.47177994f, 0.06399124f, 0.14778666f, 0.099144064f, @@ -260,11 +278,14 @@ TEST(SoftmaxOperator, ThreeDimsAxis2_opset13) { 0.23619612f, 0.1829779f, 0.37029108f, 0.14383113f, 0.0667037f, 0.15740506f, 0.13165872f, 0.31243387f, 0.24108529f, 0.15741715f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 13, /*axis*/ 2, + RunTest(input_vals_60, expected_vals, three_dimensions, /*opset*/ 13, /*axis*/ 2, + {kOpenVINOExecutionProvider}); // OpenVINO doesn't support opset-13 yet + + RunTest(input_vals_60, expected_vals, four_dimensions, /*opset*/ 13, /*axis*/ 3, {kOpenVINOExecutionProvider}); // OpenVINO doesn't support opset-13 yet } -TEST(SoftmaxOperator, ThreeDimsDefaultAxis_opset13) { +TEST(SoftmaxOperator, ThreeAndFourDimsDefaultAxis_opset13) { std::vector expected_vals = { 0.22277209f, 0.20394778f, 0.09983283f, 0.33927578f, 0.13417149f, 0.21729809f, 0.47177994f, 0.06399124f, 0.14778666f, 0.099144064f, @@ -281,11 +302,15 @@ TEST(SoftmaxOperator, ThreeDimsDefaultAxis_opset13) { 0.23619612f, 0.1829779f, 0.37029108f, 0.14383113f, 0.0667037f, 0.15740506f, 0.13165872f, 0.31243387f, 0.24108529f, 0.15741715f}; - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 13, /*default axis*/ -1, + RunTest(input_vals_60, expected_vals, three_dimensions, /*opset*/ 13, /*default axis*/ -1, + {kOpenVINOExecutionProvider}); // OpenVINO doesn't support opset-13 yet + + RunTest(input_vals_60, expected_vals, four_dimensions, /*opset*/ 13, /*default axis*/ -1, {kOpenVINOExecutionProvider}); // OpenVINO doesn't support opset-13 yet } -TEST(SoftmaxOperator, ThreeDimsNegativeAxis) { - // x = + +TEST(SoftmaxOperator, ThreeAndFourDimsNegativeAxis) { + // x = // node = onnx.helper.make_node('Softmax', inputs = ['x'], outputs = ['y'], axis = 2) // y = softmax_2d(x.reshape(12, 5)).reshape(3, 4, 5) // expect(node, inputs = [x], outputs = [y], @@ -308,7 +333,8 @@ TEST(SoftmaxOperator, ThreeDimsNegativeAxis) { 0.15740506f, 0.13165872f, 0.31243387f, 0.24108529f, 0.15741715f}; // -1 is last axis so same as axis == 2 - RunTest(x_vals_3dims, expected_vals, three_dimensions, /*opset*/ 12, /*axis*/ -1); + RunTest(input_vals_60, expected_vals, three_dimensions, /*opset*/ 12, /*axis*/ -1); + RunTest(input_vals_60, expected_vals, four_dimensions, /*opset*/ 12, /*axis*/ -1); } TEST(SoftmaxOperator, InvalidAxis) { @@ -379,5 +405,29 @@ TEST(SoftmaxOperator, 2DInputReduceOnAxis1WithLargeDim) { RunTest(x_vals, expected_vals, dimensions); } +// Regression test for NNAPI handling of a Softmax with opset < 13 where the input has been converted to NHWC. +// The NNAPI handling of the axis is different so we need to manually coerce the input to 2D, which will negate the +// layout change. Test model has a GlobalAveragePool -> Softmax which will trigger the layout change due to +// GlobalAveragePool being layout sensitive. +TEST(SoftmaxOperator, GH15949_regression_test) { + auto model_uri = ORT_TSTR("testdata/ort_github_issue_15949.onnx"); + + std::shared_ptr model; + auto status = Model::Load(model_uri, model, nullptr, GetEnvironment().GetLoggingManager()->DefaultLogger()); + + OpTester tester("Opset12NhwcSoftmax"); + tester.SetModelCache(model); + + tester.AddInput("X", {1, 3, 2, 2}, + {0.f, 1.f, 2.f, 3.f, + 4.f, 5.f, 6.f, 7.f, + 8.f, 9.f, 10.f, 11.f}); + tester.AddOutput("Y", {1, 3, 1, 1}, + {0.00032932f, 0.01798029f, 0.9816904f}); + + // disable TRT as it does not support axis=0 as used by the model + tester.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider, kCoreMLExecutionProvider}); +} + } // namespace test } // namespace onnxruntime diff --git a/onnxruntime/test/testdata/ort_github_issue_15949.onnx b/onnxruntime/test/testdata/ort_github_issue_15949.onnx new file mode 100644 index 0000000000000000000000000000000000000000..52053ddaad332cf62103ecd03cd25207f4c05c40 GIT binary patch literal 163 zcmd;J6=Eq#EiSQ|#K d69*#@GXpUb2qyt0Be6(=6fruna4`t*006&oANc?P literal 0 HcmV?d00001