[PyOV] Add Python API for MaxPool-14 and AvgPool-14

src/bindings/python/src/openvino/runtime/opset14/__init__.py

@@ -14,7 +14,7 @@
 from openvino.runtime.opset6.ops import assign
 from openvino.runtime.opset1.ops import atan
 from openvino.runtime.opset4.ops import atanh
-from openvino.runtime.opset1.ops import avg_pool
+from openvino.runtime.opset14.ops import avg_pool
 from openvino.runtime.opset5.ops import batch_norm_inference
 from openvino.runtime.opset2.ops import batch_to_space
 from openvino.runtime.opset1.ops import binary_convolution
@@ -103,7 +103,7 @@
 from openvino.runtime.opset5.ops import lstm_sequence
 from openvino.runtime.opset1.ops import matmul
 from openvino.runtime.opset8.ops import matrix_nms
-from openvino.runtime.opset8.ops import max_pool
+from openvino.runtime.opset14.ops import max_pool
 from openvino.runtime.opset1.ops import maximum
 from openvino.runtime.opset1.ops import minimum
 from openvino.runtime.opset4.ops import mish

src/bindings/python/src/openvino/runtime/opset14/ops.py

@@ -4,17 +4,20 @@
 
 """Factory functions for ops added to openvino opset14."""
 from functools import partial
-from typing import Union
+
+from typing import Union, Optional, List
 
 from openvino.runtime import Node, Type
 from openvino.runtime.opset_utils import _get_node_factory
+from openvino.runtime.utils.types import TensorShape
 from openvino.runtime.utils.decorators import nameable_op
-from openvino.runtime.utils.types import NodeInput, as_nodes
+from openvino.runtime.utils.types import NodeInput, as_node, as_nodes
 
 _get_node_factory_opset14 = partial(_get_node_factory, "opset14")
 
 
 # -------------------------------------------- ops ------------------------------------------------
+
 @nameable_op
 def convert_promote_types(
     left_node: NodeInput,
@@ -62,3 +65,103 @@ def inverse(
     }
 
     return _get_node_factory_opset14().create("Inverse", inputs, attributes)
+
+
+@nameable_op
+def max_pool(
+    data: NodeInput,
+    strides: List[int],
+    dilations: List[int],
+    pads_begin: List[int],
+    pads_end: List[int],
+    kernel_shape: TensorShape,
+    rounding_type: str = "floor",
+    auto_pad: Optional[str] = None,
+    index_element_type: Optional[str] = "i64",
+    axis: Optional[int] = 0,
+    name: Optional[str] = None,
+) -> Node:
+    """Perform max pooling operation and return both values and indices of the selected elements.
+
+    :param  data:                The node providing input data.
+    :param  strides:             The distance (in pixels) to slide the filter on the feature map
+                                 over the axes.
+    :param  dilations:           The dilation of filter elements(distance between elements).
+    :param  pads_begin:          The number of pixels to add at the beginning along each axis.
+    :param  pads_end:            The number of pixels to add at the end along each axis.
+    :param  kernel_shape:        The pooling operation kernel shape.
+    :param  rounding_type:       Determines used rounding schema when computing output shape.
+                                 Acceptable values are: ['floor', 'ceil', 'ceil_torch']. Defaults to 'floor'.
+    :param  auto_pad:            Determines how the padding is calculated. Acceptable values:
+                                 [None, 'same_upper', 'same_lower', 'valid']. Defaults to None.
+    :param  index_element_type:  The data type used for the indices output of this operator.
+                                 Defaults to i64.
+    :param  axis:                The first dimension in the data shape used to determine the maximum
+                                 returned index value. The value is the product of all dimensions
+                                 starting at the provided axis. Defaults to 0.
+    :param  name:                The optional name for the created output node.
+
+    :return:   The new node performing max pooling operation.
+    """
+    if auto_pad is None:
+        auto_pad = "explicit"
+    return _get_node_factory_opset14().create(
+        "MaxPool",
+        [as_node(data)],
+        {
+            "strides": strides,
+            "dilations": dilations,
+            "pads_begin": pads_begin,
+            "pads_end": pads_end,
+            "kernel": kernel_shape,
+            "rounding_type": rounding_type.upper(),
+            "auto_pad": auto_pad.upper(),
+            "index_element_type": index_element_type,
+            "axis": axis,
+        },
+    )
+
+
+@nameable_op
+def avg_pool(
+    data_batch: NodeInput,
+    strides: List[int],
+    pads_begin: TensorShape,
+    pads_end: TensorShape,
+    kernel_shape: TensorShape,
+    exclude_pad: bool,
+    rounding_type: str = "floor",
+    auto_pad: Optional[str] = None,
+    name: Optional[str] = None,
+) -> Node:
+    """Return average pooling node.
+
+    :param data_batch:      The input node providing data.
+    :param strides:         The window movement strides.
+    :param pads_begin:      The input data optional padding below filled with zeros.
+    :param pads_end:        The input data optional padding below filled with zeros.
+    :param kernel_shape:    The pooling window shape.
+    :param exclude_pad:     Whether or not to include zero padding in average computations.
+    :param rounding_type:   Determines used rounding schema when computing output shape. Acceptable
+                            values are: ['floor', 'ceil', 'ceil_torch']. Defaults to 'floor'.
+    :param auto_pad:        Determines how the padding is calculated. Acceptable values:
+                            [None, 'same_upper', 'same_lower', 'valid']. Defaults to None.
+    :param name:            Optional name for the new output node.
+
+    :return: New node with AvgPool operation applied on its data.
+    """
+    if auto_pad is None:
+        auto_pad = "explicit"
+    return _get_node_factory_opset14().create(
+        "AvgPool",
+        [as_node(data_batch)],
+        {
+            "strides": strides,
+            "pads_begin": pads_begin,
+            "pads_end": pads_end,
+            "kernel": kernel_shape,
+            "exclude-pad": exclude_pad,
+            "rounding_type": rounding_type.upper(),
+            "auto_pad": auto_pad.upper(),
+        },
+    )

src/bindings/python/tests/test_graph/test_pooling.py

@@ -5,7 +5,7 @@
 import numpy as np
 import pytest
 
-import openvino.runtime.opset8 as ov
+import openvino.runtime.opset14 as ov
 from openvino import Type
 
 
@@ -32,6 +32,66 @@ def test_avg_pool_2d(ndarray_1x1x4x4):
     assert node.get_output_element_type(0) == Type.f32
 
 
+def test_avg_pooling_3d_ceil_torch():
+    data = np.arange(0, 25, dtype=np.float32).reshape(1, 1, 5, 5)
+    param = ov.parameter(list(data.shape))
+    kernel_shape = [2, 2]
+    strides = [2, 2]
+    pads_begin = [1, 1]
+    pads_end = [1, 1]
+
+    node = ov.avg_pool(param, strides, pads_begin, pads_end, kernel_shape, False, "ceil_torch")
+    assert node.get_type_name() == "AvgPool"
+    assert node.get_output_size() == 1
+    assert list(node.get_output_shape(0)) == [1, 1, 3, 3]
+    assert node.get_output_element_type(0) == Type.f32
+
+
+def test_avg_pool_2d_ceil_torch():
+    input_data = np.arange(0, 243, dtype=np.float32).reshape(1, 3, 9, 9)
+    param = ov.parameter(input_data.shape, name="A", dtype=np.float32)
+    kernel_shape = [2, 2]
+    strides = [2, 2]
+    pads_begin = [1, 1]
+    pads_end = [1, 1]
+
+    node = ov.avg_pool(param, strides, pads_begin, pads_end, kernel_shape, False, "ceil_torch")
+    assert node.get_type_name() == "AvgPool"
+    assert node.get_output_size() == 1
+    assert list(node.get_output_shape(0)) == [1, 3, 5, 5]
+    assert node.get_output_element_type(0) == Type.f32
+
+
+def test_avg_pool_3x3_kernel_ceil_torch_strides():
+    input_data = np.arange(0, 300, dtype=np.float32).reshape(1, 3, 10, 10)
+    param = ov.parameter(input_data.shape, name="A", dtype=np.float32)
+    kernel_shape = [3, 3]
+    strides = [2, 2]
+    pads_begin = [0, 0]
+    pads_end = [0, 0]
+
+    node = ov.avg_pool(param, strides, pads_begin, pads_end, kernel_shape, False, "ceil_torch")
+    assert node.get_type_name() == "AvgPool"
+    assert node.get_output_size() == 1
+    assert list(node.get_output_shape(0)) == [1, 3, 5, 5]
+    assert node.get_output_element_type(0) == Type.f32
+
+
+def test_avg_pool_3x3_kernel_ceil_torch_no_strides():
+    input_data = np.arange(0, 300, dtype=np.float32).reshape(1, 3, 10, 10)
+    param = ov.parameter(input_data.shape, name="A", dtype=np.float32)
+    kernel_shape = [3, 3]
+    strides = [1, 1]
+    pads_begin = [0, 0]
+    pads_end = [0, 0]
+
+    node = ov.avg_pool(param, strides, pads_begin, pads_end, kernel_shape, False, "ceil_torch")
+    assert node.get_type_name() == "AvgPool"
+    assert node.get_output_size() == 1
+    assert list(node.get_output_shape(0)) == [1, 3, 8, 8]
+    assert node.get_output_element_type(0) == Type.f32
+
+
 def test_avg_pooling_3d(ndarray_1x1x4x4):
     data = ndarray_1x1x4x4
     data = np.broadcast_to(data, (1, 1, 4, 4, 4))
@@ -43,7 +103,7 @@ def test_avg_pooling_3d(ndarray_1x1x4x4):
     pads_end = [0] * spatial_dim_count
     exclude_pad = True
 
-    node = ov.avg_pool(param, strides, pads_begin, pads_end, kernel_shape, exclude_pad)
+    node = ov.avg_pool(param, strides, pads_begin, pads_end, kernel_shape, exclude_pad, "ceil_torch")
     assert node.get_type_name() == "AvgPool"
     assert node.get_output_size() == 1
     assert list(node.get_output_shape(0)) == [1, 1, 2, 2, 2]
@@ -265,3 +325,96 @@ def test_max_pool_same_lower_auto_pads():
     assert list(node.get_output_shape(1)) == [1, 1, 4, 4]
     assert node.get_output_element_type(0) == Type.f32
     assert node.get_output_element_type(1) == Type.i32
+
+
+def test_max_pool_strides_ceil_torch():
+    data = np.arange(0.5, 75, dtype=np.float32).reshape((1, 3, 5, 5))
+    strides = [2, 2]
+    dilations = [1, 1]
+    pads_begin = [1, 1]
+    pads_end = [1, 1]
+    kernel_shape = [2, 2]
+    rounding_type = "ceil_torch"
+    auto_pad = None
+    index_et = "i32"
+
+    data_node = ov.parameter(data.shape, name="A", dtype=np.float32)
+    node = ov.max_pool(
+        data_node,
+        strides,
+        dilations,
+        pads_begin,
+        pads_end,
+        kernel_shape,
+        rounding_type,
+        auto_pad,
+        index_et,
+    )
+    assert node.get_type_name() == "MaxPool"
+    assert node.get_output_size() == 2
+    assert list(node.get_output_shape(0)) == [1, 3, 3, 3]
+    assert list(node.get_output_shape(1)) == [1, 3, 3, 3]
+    assert node.get_output_element_type(0) == Type.f32
+    assert node.get_output_element_type(1) == Type.i32
+
+
+def test_max_pool_same_lower_auto_pads_ceil_torch():
+    data = np.arange(0.5, 16, dtype=np.float32).reshape((1, 1, 4, 4))
+    strides = [1, 1]
+    dilations = [1, 1]
+    pads_begin = [0, 0]
+    pads_end = [0, 0]
+    kernel_shape = [2, 2]
+    auto_pad = "same_lower"
+    rounding_type = "ceil_torch"
+    index_et = "i32"
+
+    data_node = ov.parameter(data.shape, name="A", dtype=np.float32)
+    node = ov.max_pool(
+        data_node,
+        strides,
+        dilations,
+        pads_begin,
+        pads_end,
+        kernel_shape,
+        rounding_type,
+        auto_pad,
+        index_et,
+    )
+    assert node.get_type_name() == "MaxPool"
+    assert node.get_output_size() == 2
+    assert list(node.get_output_shape(0)) == [1, 1, 4, 4]
+    assert list(node.get_output_shape(1)) == [1, 1, 4, 4]
+    assert node.get_output_element_type(0) == Type.f32
+    assert node.get_output_element_type(1) == Type.i32
+
+
+def test_max_pool_kernel_shape3x3_ceil_torch():
+    data = np.arange(0.5, 300, dtype=np.float32).reshape((1, 3, 10, 10))
+    strides = [1, 1]
+    dilations = [1, 1]
+    pads_begin = [0, 0]
+    pads_end = [0, 0]
+    kernel_shape = [3, 3]
+    rounding_type = "ceil_torch"
+    auto_pad = None
+    index_et = "i32"
+
+    data_node = ov.parameter(data.shape, name="A", dtype=np.float32)
+    node = ov.max_pool(
+        data_node,
+        strides,
+        dilations,
+        pads_begin,
+        pads_end,
+        kernel_shape,
+        rounding_type,
+        auto_pad,
+        index_et,
+    )
+    assert node.get_type_name() == "MaxPool"
+    assert node.get_output_size() == 2
+    assert list(node.get_output_shape(0)) == [1, 3, 8, 8]
+    assert list(node.get_output_shape(1)) == [1, 3, 8, 8]
+    assert node.get_output_element_type(0) == Type.f32
+    assert node.get_output_element_type(1) == Type.i32