Add argmax and reduction ops to transpose optimizer (#1383)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tests/test_optimizers.py

@@ -1130,6 +1130,51 @@ def test_transpose_reducesum(self, input_shape, output_shape, axes, perm_input,
         self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    @parameterized.expand([
+        ((1, 3, 4, 5), (1, 3, 4), [2], [0, 2, 3, 1], [0, 2, 1]),
+        ((1, 3, 4, 5), (1, 3), [1, 2], [0, 2, 3, 1], [0, 1]),
+        ((1, 3, 4, 5), (), [0, 1, 2, 3], [0, 2, 3, 1], []),
+        ((1, 3, 4, 5, 6), (1, 3, 5, 6), [1], [0, 2, 3, 4, 1], [0, 3, 1, 2]),
+        ((1, 3, 4, 5, 6), (1, 3), [1, 2, 3], [0, 2, 3, 4, 1], [0, 1]),
+        ((1, 3, 4, 5, 6), (), [0, 1, 2, 3, 4], [0, 2, 3, 4, 1], []),
+    ])
+    def test_transpose_reducemax(self, input_shape, output_shape, axes, perm_input, perm_output):
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=perm_input, name="trans_1")
+        node1 = helper.make_node("ReduceMax", ["Y"], ["Z"], axes=axes,
+                                 keepdims=0, name="reducemax")
+        if perm_output:
+            node2 = helper.make_node("Transpose", ["Z"], ["res"], perm=perm_output, name="trans_2")
+        else:
+            node2 = helper.make_node("Identity", ["Z"], ["res"], name="trans_2")
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-reducemax-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
+    def test_transpose_argmax(self):
+        input_shape = [1, 2, 3, 4]
+        node0 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans_1")
+        node1 = helper.make_node("ArgMax", ["Y"], ["Z"], axis=3, keepdims=0, name="argmax")
+        node2 = helper.make_node("Cast", ["Z"], ["res"], to=TensorProto.INT32, name="cast")
+
+        graph = helper.make_graph(
+            [node0, node1, node2],
+            "transpose-argmax-test",
+            [helper.make_tensor_value_info("X", TensorProto.FLOAT, input_shape)],
+            [helper.make_tensor_value_info("res", TensorProto.INT32, [1, 3, 4])],
+        )
+
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"X": np.random.randn(*input_shape).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
     @parameterized.expand([
         ((1, 3, 4, 5), (1, 3, 4, 1), [2], [0, 2, 3, 1], [0, 3, 1, 2]),
         ((1, 3, 4, 5), (1, 3, 1, 1), [1, 2], [0, 2, 3, 1], [0, 3, 1, 2]),

tf2onnx/optimizer/transpose_optimizer.py

@@ -55,7 +55,7 @@ def nodes(self):
 
     def pre_optimize_action(self):
         # make Reshape into a const, which then can be fused into Conv's weight for mobilenet_v1_75_192
-        self._output_names = [name.split(":")[0] for name in self._g.outputs]
+        self._output_names = [self._g.get_node_by_output(out).name for out in self._g.outputs]
         ops = self.nodes
         constable_reshape_ops = [n for n in ops
                                  if (n.type == "Reshape"
@@ -179,6 +179,8 @@ def _optimize_at_current_graph_level(self, graph):
     def _initialize_handlers(self):
         self._handler_map = {
             "Add": self._add_handler,
+            "ArgMax": self._arg_min_max_handler,
+            "ArgMin": self._arg_min_max_handler,
             "Cast": self._simple_through_handler,
             "Clip": self._simple_through_handler,
             "Concat": self._concat_handler,
@@ -192,8 +194,14 @@ def _initialize_handlers(self):
             "Mul": self._mul_handler,
             "Pad": self._pad_handler,
             "Reciprocal": self._simple_through_handler,
-            "ReduceMean": self._reducemean_handler,
+            "ReduceLogSum": self._reduce_handler,
+            "ReduceLogSumExp": self._reduce_handler,
+            "ReduceMax": self._reduce_handler,
+            "ReduceMean": self._reduce_handler,
+            "ReduceMin": self._reduce_handler,
+            "ReduceProd": self._reduce_handler,
             "ReduceSum": self._reducesum_handler,
+            "ReduceSumSquare": self._reduce_handler,
             "Relu": self._simple_through_handler,
             "Shape": self._shape_handler,
             "Sigmoid": self._simple_through_handler,
@@ -258,7 +266,7 @@ def _get_input_index_for_trans(self, node, trans):
         return input_index
 
     # the assumption is: both node and trans have only 1 output
-    def _switch_transpose_and_node(self, node, trans):
+    def _switch_transpose_and_node(self, node, trans, update_shape=True):
         if not self._nodes_has_single_consumer_node([trans]):
             return False
 
@@ -271,7 +279,7 @@ def _switch_transpose_and_node(self, node, trans):
         # need to transpose node shape in backward direction as well after switch
         # otherwise, reshape added in post_optimize_action may not work correctly
         shape = self._g.get_shape(node.output[0])
-        if shape:
+        if update_shape and shape:
             # only nhwc transpose can reach here
             new_shape = [shape[i] for i in NHWC_TO_NCHW]
             self._g.set_shape(node.output[0], new_shape)
@@ -700,31 +708,49 @@ def _pad_handler(self, trans, node):
         self._g.replace_input(node, node.input[1], new_pads.output[0], 1)
         return self._switch_transpose_and_node(node, trans)
 
-    def _reducemean_handler(self, trans, node):
-        axes = node.get_attr("axes").ints
-        keepdims = node.get_attr("keepdims")
+    def _arg_min_max_handler(self, trans, node):
+        axis = node.get_attr_value("axis", 0)
+        node.set_attr("axes", [axis])
+        result = self._reduce_handler(trans, node)
+        new_axis = node.get_attr_value("axes")[0]
+        node.set_attr("axis", new_axis)
+        del node.attr["axes"]
+        return result
+
+    def _reduce_handler(self, trans, node):
+        keepdims = node.get_attr_value("keepdims", 1)
         trans_rank = get_transpose_rank(trans)
-        # make sure keepdims is 1, then we can do the swap, otherwise, please don't, because
-        # once keepdims is not set, original dims are lost, so transpose back won't work well.
-        # by default, if keepdims is not specified, it is 1
-        if axes == list(range(1, trans_rank - 1)) and ((keepdims and keepdims.i == 1) or (not keepdims)):
-            node.set_attr("axes", list(range(2, trans_rank)))
-            return self._switch_transpose_and_node(node, trans)
-        return False
+        axes = node.get_attr_value("axes", list(range(trans_rank)))
+        perm = trans.get_attr("perm").ints
+        axes = [a + trans_rank if a < 0 else a for a in axes]
+        new_axes = [perm[a] for a in axes]
+        update_shape = keepdims == 1
+        shape = self._g.get_shape(node.output[0])
+        if not self._switch_transpose_and_node(node, trans, update_shape):
+            return False
+        node.set_attr("axes", new_axes)
+        if keepdims == 0:
+            remaining_axes = []
+            j = 0
+            for i in range(trans_rank):
+                if i in new_axes:
+                    remaining_axes.append(None)
+                else:
+                    remaining_axes.append(j)
+                    j += 1
+            new_perm = [remaining_axes[p] for p in perm if remaining_axes[p] is not None]
+            if shape:
+                new_shape = [shape[new_perm.index(i)] for i in range(len(new_perm))]
+                self._g.set_shape(node.output[0], new_shape)
+            trans.set_attr("perm", new_perm)
+        return True
 
     def _reducesum_handler(self, trans, node):
         keepdims = node.get_attr("keepdims")
-        # make sure keepdims is 1, then we can do the swap, otherwise, please don't, because
-        # once keepdims is not set, original dims are lost, so transpose back won't work well.
-        # by default, if keepdims is not specified, it is 1
+        if self._g.opset <= 12:
+            return self._reduce_handler(trans, node)
         if keepdims and keepdims.i == 0:
             return False
-        if self._g.opset <= 12:
-            axes = node.get_attr("axes").ints
-            perm = trans.get_attr('perm').ints
-            new_axes = [perm[axis] for axis in axes]
-            node.set_attr("axes", new_axes)
-            return self._switch_transpose_and_node(node, trans)
         if node.inputs[1].is_const():
             axes = node.inputs[1].get_tensor_value()
             perm = trans.get_attr('perm').ints