Fix some issues with NNAPI Softmax (#16095)

### Description
<!-- Describe your changes. -->
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
<!-- - Why is this change required? What problem does it solve?
- If it fixes an open issue, please link to the issue here. -->
#15949
This commit is contained in:
Scott McKay 2023-06-08 13:56:06 +10:00 committed by GitHub
parent dc1312cfb1
commit b07b647f66
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
3 changed files with 157 additions and 50 deletions

View file

@ -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<uint32_t> 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<int32_t>(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<uint32_t> 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<int32_t>(dim0), narrow<int32_t>(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<int32_t> 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<int32_t>(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<int32_t>(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<int32_t>(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;
}
}
}

View file

@ -1,12 +1,14 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <cmath>
#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 <cmath>
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<int64_t> three_dimensions = {3, 4, 5};
static std::vector<float> x_vals_3dims = {
static const std::vector<int64_t> three_dimensions = {3, 4, 5};
static const std::vector<int64_t> four_dimensions = {1, 3, 4, 5};
static const std::vector<float> 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<float> 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 = <see x_vals_3dims>
TEST(SoftmaxOperator, ThreeAndFourDimsAxis0) {
// x = <see input_vals_60>
// 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 = <see x_vals_3dims>
TEST(SoftmaxOperator, ThreeAndFourDimsSecondLastAxis) {
// x = <see input_vals_60>
// 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 = <see x_vals_3dims>
TEST(SoftmaxOperator, ThreeAndFourDimsLastAxis) {
// x = <see input_vals_60>
// 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<float> 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<float> 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 = <see x_vals_3dims>
TEST(SoftmaxOperator, ThreeAndFourDimsNegativeAxis) {
// x = <see input_vals_60>
// 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> model;
auto status = Model::Load(model_uri, model, nullptr, GetEnvironment().GetLoggingManager()->DefaultLogger());
OpTester tester("Opset12NhwcSoftmax");
tester.SetModelCache(model);
tester.AddInput<float>("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<float>("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

Binary file not shown.