diff --git a/src/common/transformations/src/transformations/common_optimizations/reverse_shape_and_type_infer.cpp b/src/common/transformations/src/transformations/common_optimizations/reverse_shape_and_type_infer.cpp index 25708ad9420..5f93afb68c0 100644 --- a/src/common/transformations/src/transformations/common_optimizations/reverse_shape_and_type_infer.cpp +++ b/src/common/transformations/src/transformations/common_optimizations/reverse_shape_and_type_infer.cpp @@ -5,6 +5,7 @@ #include "transformations/common_optimizations/reverse_shape_and_type_infer.hpp" #include "itt.hpp" +#include "openvino/core/rt_info.hpp" #include "openvino/core/validation_util.hpp" #include "openvino/op/concat.hpp" #include "openvino/op/convert_like.hpp" @@ -176,13 +177,31 @@ bool ov::pass::ReverseShapeAndTypeInfer::run_on_model(const std::shared_ptr(op)) { - auto in0_rank = op->get_input_partial_shape(0).rank(); - if (output_shape.rank().is_static() && in0_rank.is_dynamic() && op->get_input_size() > 1) { - auto in1_pshape = op->get_input_partial_shape(1); - if (in1_pshape.is_static()) { - auto num_dims = in1_pshape.size() == 0 ? 1 : in1_pshape[0].get_length(); - op->get_input_tensor(0).m_partial_shape = - PartialShape::dynamic(output_shape.rank().get_length() + num_dims); + auto in0_pshape = op->get_input_partial_shape(0); + auto in0_rank = in0_pshape.rank(); + if (output_shape.rank().is_static()) { + if (in0_rank.is_dynamic() && op->get_input_size() > 1) { + auto in1_pshape = op->get_input_partial_shape(1); + if (in1_pshape.is_static()) { + auto num_dims = in1_pshape.size() == 0 ? 1 : in1_pshape[0].get_length(); + op->get_input_tensor(0).m_partial_shape = + PartialShape::dynamic(output_shape.rank().get_length() + num_dims); + } + } else if (in0_rank.is_static() && op->get_input_size() == 1) { + // attempt to create second input + std::vector in1_data; + for (size_t i = 0; i < in0_pshape.size(); i++) { + if (in0_pshape[i] == 1) { + in1_data.push_back(i); + } + } + int64_t num_ones = in1_data.size(); + if (num_ones == in0_rank.get_length() - output_shape.rank().get_length()) { + auto axes = ov::op::v0::Constant::create(element::i64, Shape{in1_data.size()}, in1_data); + auto new_squeeze = std::make_shared(op->get_input_source_output(0), axes); + op->output(0).replace(new_squeeze->output(0)); + copy_runtime_info(op, new_squeeze); + } } } is_changed |= inherit_output_type(op, {0}); diff --git a/src/common/transformations/tests/common_optimizations/reverse_shape_and_type_infer.cpp b/src/common/transformations/tests/common_optimizations/reverse_shape_and_type_infer.cpp index f550b92ea07..82ceee31375 100644 --- a/src/common/transformations/tests/common_optimizations/reverse_shape_and_type_infer.cpp +++ b/src/common/transformations/tests/common_optimizations/reverse_shape_and_type_infer.cpp @@ -412,6 +412,41 @@ TEST_F(TransformationTestsF, SqueezeReverseInfer) { } } +TEST_F(TransformationTestsF, SqueezeAxesReverseInfer) { + auto dyn = Dimension::dynamic(); + { + auto data = std::make_shared(element::dynamic, PartialShape{1, dyn, 1, dyn, dyn, dyn}); + auto squeeze = std::make_shared(data); + // Convolution is needed to produce static rank + auto weights = + opset10::Constant::create(element::f32, Shape{64, 3, 7, 7}, std::vector(64 * 3 * 7 * 7, 0.1f)); + auto conv = std::make_shared(squeeze, + weights, + Strides{2, 2}, + CoordinateDiff{3, 3}, + CoordinateDiff{3, 3}, + Strides{1, 1}); + auto result = std::make_shared(conv); + model = std::make_shared(ResultVector{result}, ParameterVector{data}); + manager.register_pass(); + } + { + auto data = std::make_shared(element::f32, PartialShape{1, dyn, 1, dyn, dyn, dyn}); + auto axes = opset10::Constant::create(element::i64, Shape{2}, {0, 2}); + auto squeeze = std::make_shared(data, axes); + auto weights = + opset10::Constant::create(element::f32, Shape{64, 3, 7, 7}, std::vector(64 * 3 * 7 * 7, 0.1f)); + auto conv = std::make_shared(squeeze, + weights, + Strides{2, 2}, + CoordinateDiff{3, 3}, + CoordinateDiff{3, 3}, + Strides{1, 1}); + auto result = std::make_shared(conv); + model_ref = std::make_shared(ResultVector{result}, ParameterVector{data}); + } +} + TEST_F(TransformationTestsF, UnsqueezeReverseInfer) { { auto data = std::make_shared(element::dynamic, PartialShape::dynamic()); diff --git a/src/frontends/tensorflow/tests/convert_tricky_models.cpp b/src/frontends/tensorflow/tests/convert_tricky_models.cpp index f7ef57511ea..1ece7f7bec8 100644 --- a/src/frontends/tensorflow/tests/convert_tricky_models.cpp +++ b/src/frontends/tensorflow/tests/convert_tricky_models.cpp @@ -664,7 +664,8 @@ TEST_F(FrontEndConversionWithReferenceTestsF, NonMaxSuppressionWithNamedOutputs) selected_scores = make_shared(selected_scores, i32); // compute the third output - valid_outputs - Output valid_outputs = make_shared(nms->output(2)); + auto squeeze_axes = make_shared(i64, Shape{1}, 0); + Output valid_outputs = make_shared(nms->output(2), squeeze_axes); // make post-processing before the concatenation auto const_minus_one = make_shared(i32, Shape{1}, -1);