Merge pull request #535 from zhijxu-MS/fix_maxpool_bug

fix bug in maxpool

Author: nbcsm

tests/common.py

@@ -12,6 +12,7 @@
 from distutils.version import LooseVersion
 from parameterized import parameterized
 import numpy as np
+import tensorflow as tf
 from tf2onnx import constants, logging, utils
 
 __all__ = [
@@ -22,6 +23,7 @@
     "check_tf_min_version",
     "check_tf_max_version",
     "skip_tf_versions",
+    "skip_tf_cpu",
     "check_onnxruntime_min_version",
     "check_opset_min_version",
     "check_opset_max_version",
@@ -181,6 +183,15 @@ def skip_tf_versions(excluded_versions, message=""):
     return unittest.skipIf(condition, reason)
 
 
+def is_tf_gpu():
+    return tf.test.is_gpu_available()
+
+
+def skip_tf_cpu(message=""):
+    is_tf_cpu = not is_tf_gpu()
+    return unittest.skipIf(is_tf_cpu, message)
+
+
 def check_opset_min_version(min_required_version, message=""):
     """ Skip if opset < min_required_version """
     config = get_test_config()

tests/test_backend.py

@@ -258,6 +258,19 @@ def test_maxpool(self):
                 self.logger.debug(str(p))
                 self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    @skip_tf_cpu("only tf_gpu can run maxpool with NCHW format")
+    def test_maxpool_gpu(self):
+        # make sure converter behaves well when data format is NCHW
+        # and when data format is NCHW, only gpu version of tensorflow can run it.
+        ksize = [1, 1, 2, 2]
+        strides = [1, 1, 2, 2]
+        x_val = make_xval([1, 3, 50, 80])
+        for padding in ["SAME", "VALID"]:
+            x = tf.placeholder(tf.float32, shape=[None] * 4, name=_TFINPUT)
+            mp = tf.nn.max_pool(x, ksize, strides, padding=padding, data_format="NCHW")
+            _ = tf.identity(mp, name=_TFOUTPUT)
+            self._run_test_case([_OUTPUT], {_INPUT: x_val})
+
     @check_onnxruntime_incompatibility("AveragePool")
     def test_avgpool(self):
         for tf_shape in ["known", "unknown"]:

tf2onnx/graph.py

@@ -136,7 +136,7 @@ def data_format(self, val):
         self.set_attr("data_format", val)
 
     def is_nhwc(self):
-        """Return True if node is in NCHW format."""
+        """Return True if node is in NHWC format."""
         return self.data_format == "NHWC"
 
     def is_const(self):

tf2onnx/onnx_opset/nn.py

@@ -301,25 +301,26 @@ def _convert(cls, ctx, node, **kwargs):
         # T Y = MaxPool(T X, @AttrType.STRING auto_pad, @AttrType.INTS kernel_shape, @AttrType.INTS pads,
         #               @AttrType.INTS strides)
         # above seems wrong - input[1] is ksize, input[2] is strides
+        # stride and ksize in tf is not always NHWC, so watch out when converting into onnx's HCHW
         if len(node.input) < 3:
-            kernel_shape = node.get_attr("ksize").ints
-            kernel_shape = [kernel_shape[1], kernel_shape[2]]
-            node.set_attr("kernel_shape", kernel_shape)
-            strides = conv_dims_attr(node, "strides")
+            kernel_shape_tf = node.get_attr("ksize").ints
+            strides_tf = node.get_attr("strides").ints
         else:
-            kernel_shape = node.inputs[1].get_tensor_value()
-            kernel_shape = [kernel_shape[1], kernel_shape[2]]
-            node.set_attr("kernel_shape", kernel_shape)
-
-            strides = node.inputs[2].get_tensor_value()
-            strides = [strides[1], strides[2]]
-            node.set_attr("strides", strides)
-
+            kernel_shape_tf = node.inputs[1].get_tensor_value()
+            strides_tf = node.inputs[2].get_tensor_value()
             ctx.remove_input(node, node.input[2])
             ctx.remove_input(node, node.input[1])
 
+        if node.is_nhwc():
+            kernel_shape_hw = kernel_shape_tf[1:3]
+            strides_hw = strides_tf[1:3]
+        else:
+            kernel_shape_hw = kernel_shape_tf[2:4]
+            strides_hw = strides_tf[2:4]
+        node.set_attr("kernel_shape", kernel_shape_hw)
+        node.set_attr("strides", strides_hw)
         conv_dims_attr(node, "dilations")
-        add_padding(ctx, node, kernel_shape, strides)
+        add_padding(ctx, node, kernel_shape_hw, strides_hw)
         conv_convert_inputs(ctx, node, with_kernel=False)