Merge pull request #498 from chinhuang007/add-maxpoolwithargmax

Add MaxPoolWithArgmax operator support

Author: guschmue

tests/test_backend.py

@@ -87,6 +87,17 @@ def get_conv_getdata(kind=1):
     else:
         raise ValueError("kind not known")
 
+def get_maxpoolwithargmax_getdata():
+    data = [
+        ('SAME', [1, 3, 3, 1], [1, 3, 3, 1], [1, 2, 2, 1]),
+        ('SAME', [1, 5, 5, 1], [1, 4, 4, 1], [1, 2, 2, 1]),
+        ('SAME', [1, 10, 5, 1], [1, 2, 2, 1], [1, 2, 2, 1]),
+        ('SAME', [1, 10, 5, 1], [1, 4, 4, 1], [1, 1, 1, 1]),
+        ('VALID', [1, 3, 3, 1], [1, 3, 3, 1], [1, 2, 2, 1]),
+        ('VALID', [1, 5, 5, 1], [1, 4, 4, 1], [1, 2, 2, 1]),
+    ]
+    for idx, v in enumerate(data):
+        yield (idx,) + v
 
 class BackendTests(Tf2OnnxBackendTestBase):
     def _run_test_case(self, output_names_with_port, feed_dict, **kwargs):
@@ -2236,6 +2247,24 @@ def test_thresholded_relu(self):
                                 graph_validator=lambda g: check_op_count(g, "ThresholdedRelu", 1))
             tf.reset_default_graph()
 
+    @check_tf_min_version("1.13")
+    @check_opset_min_version(8, "MaxPoolWithArgmax")
+    def test_maxpoolwithargmax(self):
+        for tf_shape in ["known", "unknown"]:
+            tf.reset_default_graph()
+            for p in get_maxpoolwithargmax_getdata():
+                _, padding, x_shape, ksize, strides = p
+                tf.reset_default_graph()
+                x_val = make_xval(x_shape)
+                if tf_shape == "known":
+                    x = tf.placeholder(tf.float32, shape=x_val.shape, name=_TFINPUT)
+                else:
+                    x = tf.placeholder(tf.float32, shape=[None] * x_val.ndim, name=_TFINPUT)
+                mp = tf.nn.max_pool_with_argmax(x, ksize, strides, padding=padding)
+                _ = tf.identity(mp[0], name=_TFOUTPUT)
+                _ = tf.identity(mp[1], name=_TFOUTPUT1)
+                self.logger.debug(str(p))
+                self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: x_val})
 
 if __name__ == '__main__':
     unittest_main()

tf2onnx/onnx_opset/nn.py

@@ -327,6 +327,28 @@ def _convert(cls, ctx, node, **kwargs):
         add_padding(ctx, node, kernel_shape, strides)
         conv_convert_inputs(ctx, node, with_kernel=False)
 
+@tf_op(["MaxPoolWithArgmax"], onnx_op="MaxPool")
+class MaxPoolWithArgmaxOp:
+    @classmethod
+    def version_8(cls, ctx, node, **kwargs):
+        # T output = MaxPool(T input, @list(int) ksize, @list(int) strides, @string padding, @string data_format)
+
+        # Set kernel_shape attribute
+        kernel_shape = node.get_attr("ksize").ints
+        kernel_shape = [kernel_shape[1], kernel_shape[2]]
+        node.set_attr("kernel_shape", kernel_shape)
+
+        # Set strides attribute
+        strides = node.get_attr("strides").ints
+        strides = [strides[1], strides[2]]
+        node.set_attr("strides", strides)
+
+        # The input data_format is NHWC for TF MaxPoolWithArgmax
+        node.set_attr("data_format", "NHWC")
+
+        add_padding(ctx, node, kernel_shape, strides)
+        conv_convert_inputs(ctx, node, with_kernel=False, input_indices=[0], output_indices=[0, 1])
+
 
 @tf_op(["BiasAdd", "BiasAddV1"])
 class BiasAdd:

tf2onnx/tfonnx.py

@@ -47,7 +47,7 @@ def tflist_to_onnx(node_list, shape_override):
     ignored_attr = ["unknown_rank", "_class", "Tshape", "use_cudnn_on_gpu", "Index", "Tpaddings",
                     "TI", "Tparams", "Tindices", "Tlen", "Tdim", "dynamic_size", "Tmultiples",
                     "Tblock_shape", "Tcrops", "index_type", "Taxis", "U", "maxval",
-                    "Tout", "Tlabels", "Tindex", "element_shape"]
+                    "Tout", "Tlabels", "Tindex", "element_shape", "Targmax"]
     # some stats
     op_cnt = collections.Counter()
     attr_cnt = collections.Counter()