Merge pull request #884 from jignparm/jignparm/transpose_sum_optimizer

Adds Sum(Transpose(x1), Transpose(x2),...) optimizer.

Author: guschmue

tests/test_optimizers.py

@@ -712,6 +712,30 @@ def test_two_transposes_switch_with_mul(self):
                                              "u2": np.random.randn(1, 6, 8, 9).astype(np.float32)},
                                    model_proto, remaining_transpose_num=0)
 
+    def test_many_transposes_and_constant_switch_with_sum(self):
+        constnode = self._make_onnx_const(np.array(np.random.random((1, 8, 9, 6)), dtype=np.float32), "v4")
+        node0 = helper.make_node("Transpose", ["u1"], ["v1"], perm=[0, 2, 3, 1], name="trans_0")
+        node1 = helper.make_node("Transpose", ["u2"], ["v2"], perm=[0, 2, 3, 1], name="trans_1")
+        node11 = helper.make_node("Transpose", ["u3"], ["v3"], perm=[0, 2, 3, 1], name="trans_2")
+
+        node2 = helper.make_node("Sum", ["v1", "v2", "v3", "v4"], ["x"], name="sum_1")
+        node3 = helper.make_node("Sum", ["x", "v1"], ["y"], name="sum_2")
+        node4 = helper.make_node("Transpose", ["y"], ["res"], perm=[0, 3, 1, 2], name="trans_4")
+
+        graph = helper.make_graph(
+            [constnode, node0, node1, node11, node2, node3, node4],
+            "test-transpose-mul",
+            [helper.make_tensor_value_info("u1", TensorProto.FLOAT, (1, 6, 8, 9)),
+             helper.make_tensor_value_info("u2", TensorProto.FLOAT, (1, 6, 8, 9)),
+             helper.make_tensor_value_info("u3", TensorProto.FLOAT, (1, 6, 8, 9))],
+            [helper.make_tensor_value_info("res", TensorProto.FLOAT, (1, 6, 8, 9))],
+        )
+        model_proto = self.make_model(graph, producer_name="onnx-tests")
+        self.run_transpose_compare(["res"], {"u1": np.random.randn(1, 6, 8, 9).astype(np.float32),
+                                             "u2": np.random.randn(1, 6, 8, 9).astype(np.float32),
+                                             "u3": np.random.randn(1, 6, 8, 9).astype(np.float32)},
+                                   model_proto, remaining_transpose_num=0)
+
     # Tranpose Optimizer Tests End
 
     # Identity Optimizer Tests Start

tf2onnx/optimizer/transpose_optimizer.py

@@ -184,6 +184,8 @@ def _initialize_handlers(self):
             "ReduceMean": self._reducemean_handler,
             "Relu": self._simple_through_handler,
             "Shape": self._shape_handler,
+            "Sigmoid": self._simple_through_handler,
+            "Sum": self._sum_handler,
             "Slice": self._slice_handler,
             "Split": self._split_handler,
             "Squeeze": self._squeeze_handler,
@@ -477,6 +479,50 @@ def _mul_handler(self, trans, node):
 
         return False
 
+    def _sum_handler(self, trans, node):
+        inputs = node.inputs
+        trans_shape = self._g.get_shape(trans.output[0])
+        perm = list(trans.get_attr('perm').ints)
+        untrans_idx = [perm.index(i) for i in range(len(perm))]
+
+        # check if sum(trans(x1), trans(x2), const(x3), ...) can be switched
+        for n in inputs:
+            if n.type not in ["Transpose", "Const"]:
+                return False
+            if not self._nodes_has_single_consumer_node([n]):
+                return False
+            if n.is_const():
+                # if graph is valid, op shapes should be valid
+                # const is special case, in case of broadcasting
+                # ensure rank matches
+                n_shape = self._g.get_shape(n.output[0])
+                if len(n_shape) != len(trans_shape):
+                    return False
+            else:
+                if list(n.get_attr('perm').ints) != perm:
+                    return False
+
+        # switch to trans(sum(x1, x2, x3, ...))
+        ops = self._g.get_nodes()
+        self._g.replace_all_inputs(ops, node.output[0], trans.output[0])
+        node.input = [n.output[0] if n.is_const() else n.input[0] for n in inputs]
+        trans.input[0] = node.output[0]
+
+        # adjust shape if present
+        shape = self._g.get_shape(node.output[0])
+        if shape:
+            self._g.set_shape(node.output[0], [shape[i] for i in untrans_idx])
+
+        # update constants, remove dangling transposes
+        for n in inputs:
+            if n.is_const():
+                val = n.get_tensor_value(as_list=False)
+                new_val = np.transpose(val, untrans_idx)
+                n.set_tensor_value(new_val)
+            elif n.name != trans.name:
+                self._g.remove_node(n.name)
+        return True
+
     def _identity_handler(self, trans, node):
         if node.output[0] in node.graph.outputs:
             return False