diff --git a/onnxruntime/core/optimizer/unsqueeze_elimination.cc b/onnxruntime/core/optimizer/unsqueeze_elimination.cc index 43947a2903..86b4ce3047 100644 --- a/onnxruntime/core/optimizer/unsqueeze_elimination.cc +++ b/onnxruntime/core/optimizer/unsqueeze_elimination.cc @@ -12,16 +12,6 @@ using namespace onnxruntime::common; namespace onnxruntime { Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect& rule_effect) const { - // Get "axes" attribute. It's a required attribute so can't be null (model loading would fail if it was). - const ONNX_NAMESPACE::AttributeProto& attr = *graph_utils::GetNodeAttribute(node, "axes"); - - std::vector axes; - axes.reserve(attr.ints_size()); - for (int i = 0; i < attr.ints_size(); i++) { - axes.push_back(static_cast(attr.ints(i))); - } - - // Generate new dims. NodeArg& input_def = *node.MutableInputDefs()[0]; const auto& tensor_proto = *graph_utils::GetConstantInitializer(graph, input_def.Name()); @@ -31,8 +21,24 @@ Status UnsqueezeElimination::Apply(Graph& graph, Node& node, RewriteRuleEffect& return Status::OK(); } - std::vector new_dims(axes.size() + tensor_proto.dims().size(), 0); + std::vector axes; + if (!graph_utils::GetRepeatedNodeAttributeValues(node, "axes", axes)) { + // missing 'axes'. should have failed at model load but just in case... + return Status::OK(); + } + auto num_axes = axes.size(); + auto output_rank = num_axes + tensor_proto.dims().size(); + + // handle any negative axis values + for (auto& axis : axes) { + if (axis < 0) { + axis += output_rank; + } + } + + // Generate new dims. + std::vector new_dims(output_rank, 0); for (int64_t axis : axes) { new_dims[axis] = 1; } diff --git a/onnxruntime/test/optimizer/graph_transform_test.cc b/onnxruntime/test/optimizer/graph_transform_test.cc index 3b7cd5ef9e..9dbee78498 100644 --- a/onnxruntime/test/optimizer/graph_transform_test.cc +++ b/onnxruntime/test/optimizer/graph_transform_test.cc @@ -29,6 +29,7 @@ #include "core/util/math.h" #include "test/capturing_sink.h" #include "test/framework/test_utils.h" +#include "test/providers/provider_test_utils.h" #include "test/test_environment.h" #include "gtest/gtest.h" @@ -262,12 +263,13 @@ TEST(GraphTransformationTests, FuseConvBNNoBias) { TEST(GraphTransformationTests, FuseConvBNMulAddUnsqueeze) { std::vector test_models = {"fusion/fuse-conv-bn-mul-add-unsqueeze.onnx", + "fusion/fuse-conv-bn-mul-add-unsqueeze.negative_axes.onnx", "fusion/fuse-conv-bn-mul-add-unsqueeze-no-bias.onnx"}; for (const auto& model : test_models) { string model_uri = MODEL_FOLDER + model; std::shared_ptr p_model; - ASSERT_TRUE(Model::Load(model_uri, p_model).IsOK()); + ASSERT_STATUS_OK(Model::Load(model_uri, p_model)); Graph& graph = p_model->MainGraph(); onnxruntime::GraphTransformerManager graph_transformation_mgr{5}; diff --git a/onnxruntime/test/providers/provider_test_utils.h b/onnxruntime/test/providers/provider_test_utils.h index 1d573776c8..1c91ccba71 100644 --- a/onnxruntime/test/providers/provider_test_utils.h +++ b/onnxruntime/test/providers/provider_test_utils.h @@ -24,6 +24,20 @@ #include #include "core/util/math_cpuonly.h" +// helpers to run a function and check the status, outputting any error if it fails. +// note: wrapped in do{} while(false) so the _tmp_status variable has limited scope +#define ASSERT_STATUS_OK(function) \ + do { \ + auto _tmp_status = function; \ + ASSERT_TRUE(_tmp_status.IsOK()) << _tmp_status; \ + } while (false) + +#define EXPECT_STATUS_OK(function) \ + do { \ + auto _tmp_status = function; \ + EXPECT_TRUE(_tmp_status.IsOK()) << _tmp_status; \ + } while (false) + namespace onnxruntime { class InferenceSession; struct SessionOptions; diff --git a/onnxruntime/test/testdata/transform/fusion/fuse-conv-bn-mul-add-unsqueeze.negative_axes.onnx b/onnxruntime/test/testdata/transform/fusion/fuse-conv-bn-mul-add-unsqueeze.negative_axes.onnx new file mode 100644 index 0000000000..a2002ec79f Binary files /dev/null and b/onnxruntime/test/testdata/transform/fusion/fuse-conv-bn-mul-add-unsqueeze.negative_axes.onnx differ