Merge pull request #840 from jignparm/jignparm/optimize_transpose_multiply

Push transposes past mul ops to optimize mobilenet_v3

Author: guschmue

tests/test_optimizers.py

@@ -687,6 +687,28 @@ def test_trans_can_be_replaced_with_reshape2(self):
             self.run_transpose_compare(["Y"], {"X": np.random.randn(*input_shape_np).astype(np.float32)},
                                        model_proto, remaining_transpose_num=0)
 
+    def test_two_transposes_switch_with_mul(self):
+        const_node = self._make_onnx_const(np.array(10, dtype=np.float32), "const_10")
+        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")
+
+        node2 = helper.make_node("Mul", ["v1", "v2"], ["x"], name="mul_1")
+        node3 = helper.make_node("Mul", ["x", const_node.output[0]], ["y"], name="mul_2")
+        node4 = helper.make_node("Transpose", ["y"], ["res"], perm=[0, 3, 1, 2], name="trans_3")
+
+        graph = helper.make_graph(
+            [const_node, node0, node1, 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("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)},
+                                   model_proto, remaining_transpose_num=0)
+
     # Tranpose Optimizer Tests End
 
     # Identity Optimizer Tests Start

tf2onnx/optimizer/transpose_optimizer.py

@@ -433,7 +433,24 @@ def _mul_handler(self, trans, node):
                 multiplier_input_node = input_node
 
         # node's inputs may come from one same node. if so the multiplier_input_node may be none
-        if multiplier_input_node is None or not multiplier_input_node.is_const():
+        if multiplier_input_node is None:
+            return False
+
+        # convert  mul(trans(x), trans(y)) ->  trans(mul(x, y))
+        if multiplier_input_node.type == "Transpose":
+            if is_nhwc_transpose(multiplier_input_node):
+                if not self._nodes_has_single_consumer_node([multiplier_input_node]):
+                    return False
+                input_index = self._get_input_index_for_trans(node, multiplier_input_node)
+                if not self._switch_transpose_and_node(node, trans):
+                    return False
+
+                node.input[input_index] = multiplier_input_node.input[0]
+                self._g.remove_node(multiplier_input_node.name)
+                return True
+
+        # handle const multipliers
+        if not multiplier_input_node.is_const():
             return False
         multiplier = multiplier_input_node.get_tensor_value(as_list=False)