fix bug of tranpose opt with broadcasting op

Author: zhijxu-MS

tests/test_optimizers.py

@@ -90,6 +90,27 @@ def test_transpose_with_concat(self):
                          }
             self.run_transpose_compare(["res"], feed_dict, model_proto, remaining_transpose_num=1)
 
+    def test_transpose_with_add(self):
+        # when transpose follows with a broadcasting op
+        # reshape is needed when switching transpose with this op and op need broadcast its inputs
+        node1 = helper.make_node("Transpose", ["input_data1"], ["Y"], perm=[0, 2, 3, 1], name="trans")
+        node2 = helper.make_node("Add", ["Y", "input_data2"], ["Z"], name="add")
+
+        graph = helper.make_graph(
+            [node1, node2],
+            "transpose_with_shape",
+            [helper.make_tensor_value_info("input_data1", TensorProto.FLOAT, (2, 3, 4, 5)),
+             helper.make_tensor_value_info("input_data2", TensorProto.FLOAT, (3,)),
+             ],
+            [helper.make_tensor_value_info("Z", TensorProto.FLOAT, [2, 4, 5, 3])],
+        )
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests")
+        feed_dict = {"input_data1": np.random.randn(2, 3, 4, 5).astype(np.float32),
+                     "input_data2": np.random.randn(3).astype(np.float32),
+                     }
+        self.run_transpose_compare(["Z"], feed_dict, model_proto, remaining_transpose_num=1)
+
     def test_transpose_relu(self):
         node1 = helper.make_node("Transpose", ["X"], ["Y"], perm=[0, 2, 3, 1], name="trans_1")
         node2 = helper.make_node("Relu", ["Y"], ["Z"], name="relu")

tf2onnx/optimizer/transpose_optimizer.py

@@ -307,6 +307,14 @@ def _create_transpose_pairs_after_node(self, node):
             self._g.replace_input(consumer, node.output[0], nhwc_node.output[0])
 
     def _create_transpose_pairs_before_node(self, node):
+        def shape_after_expand(ori_shape):
+            # according to broadcasting rule to expand shape to 4D
+            utils.make_sure(ori_shape.count(-1) <= 1, "shape can contain one -1 at most")
+            ori_rank = len(ori_shape)
+            utils.make_sure(ori_rank <= 4, "ONNX only supports 4D data")
+            new_shape = [1]*(4-ori_rank) + ori_shape
+            return new_shape
+
         non_nhwc_trans_inputs = []
         for input_id, n in zip(node.input, node.inputs):
             if not is_nhwc_transpose(n):
@@ -316,7 +324,16 @@ def _create_transpose_pairs_before_node(self, node):
 
         # add Transpose(0, 3, 1, 2) and Transpose(0, 2, 3, 1) before each non_nhwc_trans_consumers
         for input_id, n in non_nhwc_trans_inputs:
-            nchw_node = self._g.make_node("Transpose", [input_id], attr={"perm": [0, 3, 1, 2]})
+            shape = self._g.get_shape(n.output[0])
+            if len(shape) == 4:
+                nchw_node = self._g.make_node("Transpose", [input_id], attr={"perm": [0, 3, 1, 2]})
+            else:
+                shape_4d = shape_after_expand(shape)
+                shape_const = self._g.make_const(utils.make_name("reshape_shape"),
+                                                 np_val=np.array(shape_4d, np.int64)).output[0]
+                reshape = self._g.make_node("Reshape", [input_id, shape_const]).output[0]
+                nchw_node = self._g.make_node("Transpose", [reshape], attr={"perm": [0, 3, 1, 2]})
+
             nhwc_node = self._g.make_node("Transpose", [nchw_node.output[0]], attr={"perm": [0, 2, 3, 1]})
             self._g.replace_input(node, input_id, nhwc_node.output[0])