[NNAPI EP] Update squeeze ops (#5946)

* [NNAPI EP] Update squeeze ops
This commit is contained in:
Guoyu Wang 2020-11-26 03:00:54 -08:00 committed by GitHub
parent d52b9aca68
commit 4afdced775
No known key found for this signature in database
GPG key ID: 4AEE18F83AFDEB23
6 changed files with 78 additions and 51 deletions

View file

@ -1,6 +1,5 @@
//
// Created by daquexian on 8/3/18.
//
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include <iostream>
#include <string>

View file

@ -1,6 +1,6 @@
//
// Created by daquexian on 5/21/18.
//
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>

View file

@ -1767,6 +1767,10 @@ Status ConcatOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const
class SqueezeOpBuilder : public BaseOpBuilder {
public:
void AddInitializersToSkip(ModelBuilder& model_builder, const Node& node) const override;
static Status AddSqueezeOp(ModelBuilder& model_builder,
const std::string& node_name,
const std::string& input, const std::string& output,
vector<int32_t> axes) ORT_MUST_USE_RESULT;
private:
Status AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node) const override ORT_MUST_USE_RESULT;
@ -1779,6 +1783,49 @@ void SqueezeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const
}
}
/* static */ Status SqueezeOpBuilder::AddSqueezeOp(ModelBuilder& model_builder,
const std::string& node_name,
const std::string& input, const std::string& output,
vector<int32_t> axes) {
auto& shaper(model_builder.GetShaper());
const auto& operand_indices(model_builder.GetOperandIndices());
const auto& operand_types(model_builder.GetOperandTypes());
const auto& input_shape(shaper[input]);
auto input_dims = input_shape.size();
for (auto& axis : axes) {
axis = static_cast<int32_t>(HandleNegativeAxis(axis, input_dims));
}
// Despite the spec of ANEURALNETWORKS_SQUEEZE at
// https://developer.android.com/ndk/reference/group/neural-networks
// states, that the axes (input 1 of ANEURALNETWORKS_SQUEEZE) is optional.
//
// The actual code of NNAPI requires the axes to be provided
// https://android.googlesource.com/platform/frameworks/ml/+/master/nn/common/operations/Squeeze.cpp#31
if (axes.empty()) { // Squeeze all
for (size_t i = 0; i < input_dims; i++) {
if (input_shape[i] == 1)
axes.push_back(i);
}
}
const auto axes_name = model_builder.GetUniqueName(node_name + input + "_axes");
Shape axes_dimen = {static_cast<uint32_t>(axes.size())};
const OperandType axes_operand_type(Type::TENSOR_INT32, axes_dimen);
ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(axes_name, axes.data(), axes_operand_type));
std::vector<uint32_t> input_indices;
input_indices.push_back(operand_indices.at(input)); // input
input_indices.push_back(operand_indices.at(axes_name)); // axes
ORT_RETURN_IF_ERROR(shaper.Squeeze(input, axes, output));
const OperandType output_operand_type(operand_types.at(input).type, shaper[output]);
ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SQUEEZE, input_indices,
{output}, {output_operand_type}, {false}));
return Status::OK();
}
/* static */ vector<int32_t> SqueezeOpBuilder::GetAxes(ModelBuilder& model_builder, const Node& node) {
vector<int32_t> axes;
// Squeeze opset 13 use input as axes
@ -1804,47 +1851,13 @@ void SqueezeOpBuilder::AddInitializersToSkip(ModelBuilder& model_builder, const
}
Status SqueezeOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder, const Node& node) const {
auto& shaper(model_builder.GetShaper());
const auto& operand_indices(model_builder.GetOperandIndices());
const auto& operand_types(model_builder.GetOperandTypes());
auto input = node.InputDefs()[0]->Name();
if (model_builder.IsOperandNHWC(input)) {
// We want to transpose nhwc operand back to nchw before squeeze
ORT_RETURN_IF_ERROR(GetNCHWInput(model_builder, node, 0, input));
}
NodeAttrHelper helper(node);
vector<int32_t> axes = GetAxes(model_builder, node);
const auto& input_shape(shaper[input]);
auto input_dims = input_shape.size();
for (auto& axis : axes) {
axis = static_cast<int32_t>(HandleNegativeAxis(axis, input_dims));
}
if (axes.empty()) { // Squeeze all
for (size_t i = 0; i < input_dims; i++) {
if (input_shape[i] == 1)
axes.push_back(i);
}
}
const auto axes_name = model_builder.GetUniqueName(node.Name() + input + "_axes");
Shape axes_dimen = {static_cast<uint32_t>(axes.size())};
shaper.AddShape(axes_name, axes_dimen);
const OperandType axes_operand_type(Type::TENSOR_INT32, axes_dimen);
ORT_RETURN_IF_ERROR(model_builder.AddOperandFromPersistMemoryBuffer(axes_name, axes.data(), axes_operand_type));
std::vector<uint32_t> input_indices;
input_indices.push_back(operand_indices.at(input)); // input
input_indices.push_back(operand_indices.at(axes_name)); // axes
const auto& output = node.OutputDefs()[0]->Name();
ORT_RETURN_IF_ERROR(shaper.Squeeze(input, axes, output));
const OperandType output_operand_type(operand_types.at(input).type, shaper[output]);
ORT_RETURN_IF_ERROR(model_builder.AddOperation(ANEURALNETWORKS_SQUEEZE, input_indices,
{output}, {output_operand_type}, {false}));
return Status::OK();
return AddSqueezeOp(model_builder, node.Name(), input, node.OutputDefs()[0]->Name(), GetAxes(model_builder, node));
}
#pragma endregion

View file

@ -1,3 +1,6 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#include "core/providers/common.h"
#include "core/providers/nnapi/nnapi_builtin/nnapi_lib/NeuralNetworksWrapper.h"
@ -374,17 +377,12 @@ Status Shaper::SqueezeImpl(const std::string& input_name,
const std::string& output_name) {
const Shape& input_dimen = shape_map_.at(input_name);
int32_t input_size = input_dimen.size();
size_t axes_size = axes.size();
std::unordered_set<int32_t> axes_to_be_squeezed;
if (axes_size == 0) {
for (int32_t idx = 0; idx < input_size; ++idx) {
if (input_dimen[idx] == 1)
axes_to_be_squeezed.insert(idx);
}
} else {
for (const auto& axis : axes)
axes_to_be_squeezed.insert(axis);
}
// If the Op is squeezing all by not specifying axes, the axes is pre-populate
// with axes of all single dimensions by the caller
for (const auto& axis : axes)
axes_to_be_squeezed.insert(axis);
// Make output dimensions
std::vector<uint32_t> output_dimen;
@ -394,6 +392,11 @@ Status Shaper::SqueezeImpl(const std::string& input_name,
output_dimen.push_back(input_dimen[i]);
}
// In case of a tensor has all 1's in dimension such as {1,1,1,1} and gets squeezed all
// the output shape will be {1}
if (output_dimen.empty())
output_dimen.push_back(1);
shape_map_[output_name] = output_dimen;
return Status::OK();
}

View file

@ -1,3 +1,6 @@
// Copyright (c) Microsoft Corporation. All rights reserved.
// Licensed under the MIT License.
#pragma once
#include <string>

View file

@ -35,6 +35,15 @@ TEST(SqueezeOpTest, Squeeze_Empty_Axes_2) {
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
TEST(SqueezeOpTest, Squeeze_Empty_Axes_3) {
OpTester test("Squeeze");
// Squeeze all for all 1's shape will end up as a scalar
test.AddInput<float>("data", {1, 1, 1, 1}, std::vector<float>{1.0f});
test.AddOutput<float>("squeezed", {}, std::vector<float>{1.0f});
// TensorRT doesn't seem to support missing 'axes'
test.Run(OpTester::ExpectResult::kExpectSuccess, "", {kTensorrtExecutionProvider});
}
TEST(SqueezeOpTest, Squeeze_1_int32) {
OpTester test("Squeeze");
test.AddAttribute("axes", std::vector<int64_t>{0});