diff --git a/src/frontends/pytorch/src/op/max_poolnd.cpp b/src/frontends/pytorch/src/op/max_poolnd.cpp index b1e2b5b892f28b..0d685f03e16019 100644 --- a/src/frontends/pytorch/src/op/max_poolnd.cpp +++ b/src/frontends/pytorch/src/op/max_poolnd.cpp @@ -3,7 +3,18 @@ // #include "openvino/frontend/pytorch/node_context.hpp" +#include "openvino/op/add.hpp" +#include "openvino/op/concat.hpp" +#include "openvino/op/constant.hpp" +#include "openvino/op/gather.hpp" +#include "openvino/op/greater.hpp" #include "openvino/op/max_pool.hpp" +#include "openvino/op/multiply.hpp" +#include "openvino/op/pad.hpp" +#include "openvino/op/range.hpp" +#include "openvino/op/select.hpp" +#include "openvino/op/shape_of.hpp" +#include "openvino/op/subtract.hpp" #include "openvino/op/util/framework_node.hpp" #include "utils.hpp" @@ -21,7 +32,8 @@ OutputVector translate_max_poolnd(const NodeContext& context) { if (!context.input_is_none(2)) { strides = context.const_input(2); } - if (context.input_is_none(2) || strides.size() == 0) { + const bool use_kernel = context.input_is_none(2) || (strides.size() == 0); + if (use_kernel) { // In case strides are not provided default is kernel strides = kernel; } @@ -42,8 +54,54 @@ OutputVector translate_max_poolnd(const NodeContext& context) { rounding_type = context.const_input(5) ? RoundingType::CEIL : RoundingType::FLOOR; } - return {context.mark_node( - std::make_shared(context.get_input(0), strides, dilations, pads, pads, kernel, rounding_type))}; + auto input = context.get_input(0); + if (rounding_type == RoundingType::CEIL) { + // The corner case of Max Pooling with ceil_mode on + // PyTorch allows sliding window go off bound, which leads to this accommodation. + // More detail on https://github.com/pytorch/pytorch/issues/57178 + const auto zero = context.mark_node(v0::Constant::create(element::i32, Shape{}, {0})); + const auto one = context.mark_node(v0::Constant::create(element::i32, Shape{}, {1})); + const auto two = context.mark_node(v0::Constant::create(element::i32, Shape{}, {2})); + + const auto padding = + context.input_is_none(3) + ? context.mark_node(std::make_shared(element::i32, Shape{pads.size()}, 0))->output(0) + : context.get_input(3); + const auto pads_len = context.mark_node(v0::Constant::create(element::i32, Shape{}, {pads.size()})); + const auto pads_remaining = context.mark_node(v0::Constant::create(element::i32, Shape{2}, {0, 0})); + + // gather input spatial dims and prepare for compare as values (in_dim + pad) + const auto input_shape_rank = get_shape_rank(context, input); + const auto end = context.mark_node(v0::Constant::create(element::i32, Shape{}, {pads.size() + 2})); + const auto dim_idxs = context.mark_node(std::make_shared(two, end, one, element::i32)); + const auto gth_in_dims = + context.mark_node(std::make_shared(std::get<0>(input_shape_rank), dim_idxs, zero)); + const auto in_left_padded = context.mark_node(std::make_shared(gth_in_dims, padding)); + + // gather output spatial dims and prepare it for compare as values (out_dim - 1) * stride + const auto mp = context.mark_node( + std::make_shared(input, strides, dilations, pads, pads, kernel, rounding_type)); + const auto shape_of_mp = context.mark_node(std::make_shared(mp, element::i32)); + const auto gth_out_dims = context.mark_node(std::make_shared(shape_of_mp, dim_idxs, zero)); + const auto out_sub_one = context.mark_node(std::make_shared(gth_out_dims, one)); + const auto stride_node = use_kernel ? context.get_input(1) : context.get_input(2); + const auto out_mul_stride = context.mark_node(std::make_shared(out_sub_one, stride_node)); + + // if (in_dim + pad) > ((out_dim - 1) * stride) sliding window in bound use end padding. + const auto in_gt_out = context.mark_node(std::make_shared(in_left_padded, out_mul_stride)); + const auto selected_pads = context.mark_node(std::make_shared(in_gt_out, padding, zero)); + + // apply padding on input clear pads attribute + const auto pb = context.mark_node(std::make_shared(OutputVector{pads_remaining, padding}, 0)); + const auto pe = context.mark_node(std::make_shared(OutputVector{pads_remaining, selected_pads}, 0)); + const auto minus_inf = + context.mark_node(v0::Constant::create(element::f32, Shape{}, {-std::numeric_limits::infinity()})); + input = context.mark_node(std::make_shared(input, pb, pe, minus_inf, op::PadMode::CONSTANT)); + std::fill_n(pads.begin(), pads.size(), 0); + } + + return { + context.mark_node(std::make_shared(input, strides, dilations, pads, pads, kernel, rounding_type))}; }; OutputVector translate_max_poolnd_fx(const NodeContext& context) { diff --git a/tests/layer_tests/pytorch_tests/test_pooling.py b/tests/layer_tests/pytorch_tests/test_pooling.py index 3f3b5bab201113..bd5e7b5345396c 100644 --- a/tests/layer_tests/pytorch_tests/test_pooling.py +++ b/tests/layer_tests/pytorch_tests/test_pooling.py @@ -15,6 +15,8 @@ {'kernel_size': [2, 1], 'stride': [], 'padding': 0}, ] +d2_params_corner_case = [{'kernel_size': [8, 8], 'stride': [8,4], 'padding': 1}] + d1_params = [{'kernel_size': 3, 'stride': 1, 'padding': 0}, {'kernel_size': (4,), 'stride': 1, 'padding': 1}, {'kernel_size': 4, 'stride': (5,), 'padding': 2}, @@ -117,7 +119,15 @@ def test_avg_pool1d(self, params, ceil_mode, count_include_pad, ie_device, preci ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 3}, trace_model=True, dynamic_shapes=False) - @pytest.mark.parametrize("params", d2_params) + @pytest.mark.parametrize( + "params", + d2_params + + [ + pytest.param( + {"kernel_size": [8, 8], "stride": [8, 4], "padding": 1}, + marks=pytest.mark.xfail(reason="Sliding windows that would start in the right padded are ignored.") + ) + ]) @pytest.mark.parametrize("ceil_mode", [True, False]) @pytest.mark.parametrize("count_include_pad", [True, False]) @pytest.mark.nightly @@ -145,7 +155,7 @@ def test_max_pool1d(self, params, ceil_mode, dilation, ie_device, precision, ir_ self._test(*self.create_model("max_pool1d", **params, ceil_mode=ceil_mode, dilation=dilation), ie_device, precision, ir_version, kwargs_to_prepare_input={'ndim': 3}, dynamic_shapes=False) - @pytest.mark.parametrize("params", d2_params) + @pytest.mark.parametrize("params", d2_params + d2_params_corner_case) @pytest.mark.parametrize("ceil_mode", [True, False]) @pytest.mark.parametrize("dilation", [1, 2]) @pytest.mark.nightly