LpPool operator implementation (#685)

* LpPool operator implementation

Implements LpPool by utilizing avg_pool:

    window_size = np.prod(ksize)

    input = tf.math.pow(input, self.p) * window_size
    pooled = tf.nn.avg_pool(input, ksize=ksize, strides=strides, padding=padding)
    pooled = tf.math.pow(pooled, 1.0 / self.p)

* Fixed the test cases and few other issues with the code

* fixed the test cases to run correct lp pooling
* bug in the caculations - added abs to the formula
* fixed a problem with the SAME_UPPER padding caculating wrong values

Co-authored-by: Chin Huang <chhuang@us.ibm.com>

Author: sdmonov

doc/support_status.md

@@ -87,7 +87,7 @@ Notes:
 |LogSoftmax|**1**|1|1|1|1|1|1|1|1|1|**11**|11|**13**:small_red_triangle:|
 |Loop|**1**|1|1|1|1|1|1|1|1|1|**11**|11|11|
 |LpNormalization|**1**|1|1|1|1|1|1|1|1|1|1|1|1|
-|LpPool|**1**:small_red_triangle:|**2**:small_red_triangle:|2:small_red_triangle:|2:small_red_triangle:|2:small_red_triangle:|2:small_red_triangle:|2:small_red_triangle:|2:small_red_triangle:|2:small_red_triangle:|2:small_red_triangle:|**11**:small_red_triangle:|11:small_red_triangle:|11:small_red_triangle:|
+|LpPool|**1**|**2**|2|2|2|2|2|2|2|2|**11**|11|11|
 |MatMul|**1**|1|1|1|1|1|1|1|**9**|9|9|9|**13**:small_red_triangle:|
 |MatMulInteger|-|-|-|-|-|-|-|-|-|**10**|10|10|10|
 |Max|**1**|1|1|1|1|**6**|6|**8**|8|8|8|**12**|**13**|

onnx_tf/common/pooling_helper.py

@@ -120,7 +120,7 @@ def _pooling_output_shape(input_size, ksize, stride, dilation, pad, ceil_mode):
 
 def py_pool(input, kernel_shape, strides=None, dilations=None,
             padding=None, ceil_mode=False, pooling_type="MAX",
-            include_indices=True):
+            include_indices=True, p=2):
     """
         Implementation of Max and Average pool operations in Python
         Args:
@@ -133,8 +133,10 @@ def py_pool(input, kernel_shape, strides=None, dilations=None,
                           [x1_begin, x2_begin...x1_end, x2_end,...]
             ceil_mode:    whether to use ceil or floor (default) to compute
                           the output shape.
-            pooling_type: specify pooling type. Values can be "MAX" or "AVG".
+            pooling_type: specifies pooling type. Values can be "MAX", "AVG" or
+                          "LP"
             include_indices: should indices be included in the output
+            p:            specifies the p parameter for LpPooling
       Return:
             pooled:       output data from max pooling across the input
             ind:          indices of the selected max values from the input
@@ -151,6 +153,9 @@ def py_pool(input, kernel_shape, strides=None, dilations=None,
     else:
         input_dtype_min = np.finfo(input_dtype).min
 
+    if pooling_type == "LP":
+      rootN = (1.0 / p)
+
     def _loop_over_output(batch, channel):
         dims = [range(output_sp_shape[d]) for d in range(spatial_size)]
         for counters in itertools.product(*dims):
@@ -167,7 +172,7 @@ def _loop_over_output(batch, channel):
                 cur_range = [i for i in range(dim_start,
                                               dim_end, dilations[dim])]
                 input_ranges.append(cur_range)
-            if pooling_type == "AVG":
+            if pooling_type in ["AVG", "LP"]:
                 val_sum = 0
                 val_count = 0
             else:
@@ -179,6 +184,8 @@ def _loop_over_output(batch, channel):
                 if pooling_type == "AVG":
                     val_sum += val
                     val_count += 1
+                elif pooling_type == "LP":
+                    val_sum += abs(val ** p)
                 else:
                     if val > maxval:
                         maxval = val
@@ -192,6 +199,8 @@ def _loop_over_output(batch, channel):
             ind = (batch, channel) + counters
             if pooling_type == "AVG":
                 out_pool[ind] = val_sum / val_count
+            elif pooling_type == "LP":
+                out_pool[ind] = val_sum ** rootN
             else:
                 out_pool[ind] = maxval
                 out_ind[ind] = maxind

onnx_tf/handlers/backend/dilated_pooling.py

@@ -151,7 +151,8 @@ def __init__(self,
                padding="VALID",
                ceil_mode=False,
                count_include_pad=False,
-               pooling_type="MAX"):
+               pooling_type="MAX",
+               p=2):
     self.input = tf.convert_to_tensor(input)
 
     self.kernel_shape = kernel_shape
@@ -162,6 +163,7 @@ def __init__(self,
     self.ceil_mode = ceil_mode
     self.count_include_pad = count_include_pad
     self.pooling_type = pooling_type.upper()
+    self.p = p
 
     self.is_known_shape = self.input.shape.is_fully_defined()
     self.spatial_size = len(kernel_shape)
@@ -572,6 +574,15 @@ def dilated_maxpool_with_argmax(self, force_custom_impl=False):
 
     return (pooled, new_ind)
 
+  def _lp_pool(self, input, ksize, strides, padding):
+    window_size = np.prod(ksize)
+
+    input = tf.math.pow(tf.math.abs(input), self.p) * window_size
+    pooled = tf.nn.avg_pool(input, ksize=ksize, strides=strides, padding=padding)
+    pooled = tf.math.pow(pooled, 1.0 / self.p)
+
+    return pooled
+
   def dilated_pool(self, force_custom_impl=False):
     """
             Does N-D dilated max/avg pooling. Pads the input if explicit or
@@ -582,7 +593,7 @@ def dilated_pool(self, force_custom_impl=False):
 
     if self.is_explicit_padding or self.padding.lower() == "same_lower" \
             or (self.padding.lower() == "same_upper" and
-                self.count_include_pad):
+                self.count_include_pad) or self.pooling_type.upper() == "LP":
       # pad the input
       self._pad_input()
 
@@ -614,8 +625,8 @@ def dilated_pool(self, force_custom_impl=False):
     elif self.spatial_size < 4 and (self.strides == [1] * self.spatial_size or
             self.dilations == [1] * self.spatial_size) and \
             not force_custom_impl:
-      # if strides == 1 use tf.nn.pool directly
-      if self.strides == [1] * self.spatial_size:
+      # if strides == 1 and not LpPool use tf.nn.pool directly
+      if self.strides == [1] * self.spatial_size and self.pooling_type != "LP":
         pooled = tf.nn.pool(
             self.input,
             window_shape=self.kernel_shape,
@@ -629,6 +640,8 @@ def dilated_pool(self, force_custom_impl=False):
           op = tf.nn.max_pool
         elif self.pooling_type == "AVG":
           op = tf.nn.avg_pool
+        elif self.pooling_type == "LP":
+          op = self._lp_pool
         else:
           raise ValueError("%d-D %s pooling is not supported." %
                            (self.spatial_size, self.pooling_type))
@@ -644,12 +657,20 @@ def dilated_pool(self, force_custom_impl=False):
         # pad the input
         self._pad_input()
       input_ = self._remove_dilations()
-      pooled = tf.nn.pool(
-          input_,
-          window_shape=self.kernel_shape,
-          strides=self.kernel_shape,
-          padding="VALID",
-          pooling_type=self.pooling_type)
+      if self.pooling_type=="LP":
+        pooled = self._lp_pool(
+            input_,
+            ksize=self.kernel_shape,
+            strides=self.kernel_shape,
+            padding="VALID")
+
+      else:
+        pooled = tf.nn.pool(
+            input_,
+            window_shape=self.kernel_shape,
+            strides=self.kernel_shape,
+            padding="VALID",
+            pooling_type=self.pooling_type)
     return pooled
 
   def is_supported(self):
@@ -658,7 +679,8 @@ def is_supported(self):
             supported for average pool
         """
     # check for maxpool
-    if self.pooling_type.startswith("MAX"):
+    if self.pooling_type.startswith("MAX") or \
+       self.pooling_type=="LP":
       return True
     else:
       # if count_include_pad is true it is fully supported

onnx_tf/handlers/backend/lp_pool.py

@@ -0,0 +1,24 @@
+from onnx_tf.handlers.backend_handler import BackendHandler
+from onnx_tf.handlers.handler import onnx_op
+from .pool_mixin import PoolMixin
+
+
+@onnx_op("LpPool")
+class LpPool(PoolMixin, BackendHandler):
+
+  @classmethod
+  def _common(cls, node, **kwargs):
+    return cls.pool(node, kwargs["tensor_dict"], "LP",
+                    kwargs.get("strict", True))
+
+  @classmethod
+  def version_1(cls, node, **kwargs):
+    return cls._common(node, **kwargs)
+
+  @classmethod
+  def version_2(cls, node, **kwargs):
+    return cls._common(node, **kwargs)
+
+  @classmethod
+  def version_11(cls, node, **kwargs):
+    return cls._common(node, **kwargs)

onnx_tf/handlers/backend/pool_mixin.py

@@ -26,6 +26,8 @@ def pool(cls, node, input_dict, pooling_type, strict=True):
     dilations = node.attrs.get("dilations", [1] * spatial_size)
     ceil_mode = bool(node.attrs.get("ceil_mode", 0))
     pads = node.attrs.get("auto_pad", "NOTSET")
+    p = node.attrs.get("p", 2)
+
     if pads == "NOTSET":
       pads = node.attrs.get("pads", [0] * spatial_size * 2)
       # In case shape is fully defined, check if pads match
@@ -44,6 +46,8 @@ def pool(cls, node, input_dict, pooling_type, strict=True):
       pooling_name = "MaxPool"
     elif pooling_type == "MAX_WITH_ARGMAX":
       pooling_name = "MaxPoolWithArgmax"
+    elif pooling_type == "LP":
+      pooling_name = "LpPool"
 
     if spatial_size > 3:
       exception.OP_UNSUPPORTED_EXCEPT(
@@ -71,16 +75,16 @@ def pool(cls, node, input_dict, pooling_type, strict=True):
         padding=pads,
         ceil_mode=ceil_mode,
         pooling_type=pooling_type,
-        count_include_pad=count_include_pad)
+        count_include_pad=count_include_pad,
+        p=p)
     if not dp.is_supported():
       if strict:
         logger.warning("Using the pooling op in compatibility mode. "
-                      "This means your graph cannot be serialized.",
-                      UserWarning)
+                      "This means your graph cannot be serialized.")
 
         result = tf.numpy_function(py_pool, [
                 orig_x, kernel_shape, strides, dilations, pads, ceil_mode,
-                "AVG", False
+                pooling_type, False
             ], orig_x.dtype)
 
         if orig_x.shape.is_fully_defined():
@@ -92,15 +96,15 @@ def pool(cls, node, input_dict, pooling_type, strict=True):
         result.set_shape(output_shape)
         return [result]
       else:
-        exception.OP_UNSUPPORTED_EXCEPT("strict == 0 and average pool"
+        exception.OP_UNSUPPORTED_EXCEPT("strict == 0 and " + pooling_name +
                                         " arguments not compatible",
                                         "Tensorflow")
 
     def dilated_pool():
       return (dp.dilated_pool(), None)
 
     # select correct op depending on the pooling type
-    pooling_op = dilated_pool if pooling_type in ["MAX", "AVG"] else \
+    pooling_op = dilated_pool if pooling_type in ["MAX", "AVG", "LP"] else \
         dp.dilated_maxpool_with_argmax
 
     # select the correct transpose ops depending on the input storage format

onnx_tf/opset_version.py

@@ -86,7 +86,7 @@
     'LogSoftmax': [1, 11],
     'Loop': [1, 11],
     'LpNormalization': [1],
-    'LpPool': [],
+    'LpPool': [1, 2, 11],
     'MatMul': [1, 9],
     'MatMulInteger': [10],
     'Max': [1, 6, 8, 12, 13],