Skip to content

Commit

Permalink
[PyOV] Add Python API for MaxPool-14 and AvgPool-14 (#22966)
Browse files Browse the repository at this point in the history
### Details:
 - Extend Python API with`MaxPool-14` and `AvgPool-14`
- They both introduce a new ceil mode:
`ov::op::RoundingType::CEIL_TORCH`
- The new ceiling mode does not allow the last pooling in a Dimension to
start in the padding area

### Related PRs
 - #22930
 - #22796
 - #23027
 - #23381
 - #23582

### Tickets:
 - 131961

### Context
#18731

---------

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>
  • Loading branch information
p-wysocki and mitruska authored Mar 22, 2024
1 parent e4c7194 commit 8872077
Show file tree
Hide file tree
Showing 3 changed files with 322 additions and 252 deletions.
4 changes: 2 additions & 2 deletions src/bindings/python/src/openvino/runtime/opset14/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
107 changes: 105 additions & 2 deletions src/bindings/python/src/openvino/runtime/opset14/ops.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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[Union[str, Type]] = "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 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 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(),
},
)
Loading

0 comments on commit 8872077

Please sign in to comment.