Merge pull request #607 from zhijxu-MS/refactor

refatcor and fix bug

Author: zhijxu-MS

tests/test_optimizers.py

@@ -63,6 +63,33 @@ def check_transpose_perm(self, model_proto, expected_perm):
                 perm = list(node.attribute[0].ints)
                 self.assertEqual(perm, expected_perm)
 
+    def test_transpose_with_concat(self):
+        input_shape = (2, 3, 4, 5)
+        perm = [0, 3, 1, 2]
+        input_shape_with_trans = [input_shape[i] for i in perm]
+        for axis in [0, 1, 2, 3]:
+            output_before_trans = list(input_shape)
+            output_before_trans[axis] *= 2
+            output_shape = [output_before_trans[i] for i in [0, 3, 1, 2]]
+            node1 = helper.make_node("Transpose", ["input_data1"], ["Y"], perm=[0, 2, 3, 1], name="trans")
+            node2 = helper.make_node("Concat", ["Y", "input_data2"], ["Z"], axis=axis, name="concat")
+            node3 = helper.make_node("Transpose", ["Z"], ["res"], perm=[0, 3, 1, 2], name="trans2")
+
+            graph = helper.make_graph(
+                [node1, node2, node3],
+                "test_transpose_with_concat",
+                [helper.make_tensor_value_info("input_data1", TensorProto.FLOAT, input_shape_with_trans),
+                 helper.make_tensor_value_info("input_data2", TensorProto.FLOAT, input_shape),
+                 ],
+                [helper.make_tensor_value_info("res", TensorProto.FLOAT, output_shape)],
+            )
+
+            model_proto = helper.make_model(graph, producer_name="onnx-tests")
+            feed_dict = {"input_data1": np.random.randn(*input_shape_with_trans).astype(np.float32),
+                         "input_data2": np.random.randn(*input_shape).astype(np.float32),
+                         }
+            self.run_transpose_compare(["res"], 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/onnx_opset/nn.py

@@ -538,13 +538,13 @@ def _convert_since_9(cls, ctx, node, op_type):
             # scales is nchw
             scales = ctx.make_node("Concat", [const_one_array.output[0], scales_hw.output[0]], {"axis": 0})
         # because onnxruntime only supports to scale the last two dims so transpose is inserted
-        input_nchw = ctx.make_node("Transpose", [node.input[0]], {"perm": [0, 3, 1, 2]})
+        input_nchw = ctx.make_node("Transpose", [node.input[0]], {"perm": constants.NHWC_TO_NCHW})
         upsample = ctx.make_node(op_type, [input_nchw.output[0], scales.output[0]], attr={"mode": mode})
 
         shapes = node.output_shapes
         dtypes = node.output_dtypes
         ctx.remove_node(node.name)
-        ctx.make_node("Transpose", upsample.output, {"perm": [0, 2, 3, 1]},
+        ctx.make_node("Transpose", upsample.output, {"perm": constants.NCHW_TO_NHWC},
                       name=node.name, outputs=node.output, shapes=shapes, dtypes=dtypes)
 
 

tf2onnx/optimizer/transpose_optimizer.py

@@ -8,6 +8,7 @@
 
 import numpy as np
 
+from tf2onnx.constants import NCHW_TO_NHWC, NHWC_TO_NCHW
 from .. import utils
 from .optimizer_base import GraphOptimizerBase
 
@@ -18,12 +19,12 @@
 
 def is_nhwc_transpose(transpose_node):
     perm_attr = transpose_node.get_attr('perm')
-    return transpose_node.type == "Transpose" and perm_attr and perm_attr.ints == [0, 2, 3, 1]
+    return transpose_node.type == "Transpose" and perm_attr and perm_attr.ints == NCHW_TO_NHWC
 
 
 def is_nchw_transpose(transpose_node):
     perm_attr = transpose_node.get_attr('perm')
-    return transpose_node.type == "Transpose" and perm_attr and perm_attr.ints == [0, 3, 1, 2]
+    return transpose_node.type == "Transpose" and perm_attr and perm_attr.ints == NHWC_TO_NCHW
 
 
 def is_useless_transpose(transpose_node):
@@ -86,9 +87,9 @@ def _calculate_new_shape(graph, op):
             # reshape requires tha output shape can only contain one -1, if not some extra op needed.
             input_shape = graph.make_node("Shape", [op.input[0]]).output[0]
             if is_nchw_transpose(op):
-                indice = graph.make_const(utils.make_name("indice"), np.array([0, 3, 1, 2])).output[0]
+                indice = graph.make_const(utils.make_name("indice"), np.array(NHWC_TO_NCHW)).output[0]
             else:
-                indice = graph.make_const(utils.make_name("indice"), np.array([0, 2, 3, 1])).output[0]
+                indice = graph.make_const(utils.make_name("indice"), np.array(NCHW_TO_NHWC)).output[0]
 
             return graph.make_node("Gather", [input_shape, indice]).output[0]
 
@@ -245,7 +246,7 @@ def _switch_transpose_and_node(self, node, trans):
         shape = self._g.get_shape(node.output[0])
         if shape:
             # only nhwc transpose can reach here
-            new_shape = [shape[i] for i in [0, 3, 1, 2]]
+            new_shape = [shape[i] for i in NHWC_TO_NCHW]
             self._g.set_shape(node.output[0], new_shape)
         return True
 
@@ -301,8 +302,8 @@ def _create_transpose_pairs_after_node(self, node):
         non_nchw_trans_consumers = self._get_non_nchw_transpose_output_nodes(node)
         # add Transpose(0, 3, 1, 2) and Transpose(0, 2, 3, 1) before each non_nchw_trans_consumers
         for consumer in non_nchw_trans_consumers:
-            nchw_node = self._g.make_node("Transpose", [node.output[0]], attr={"perm": [0, 3, 1, 2]})
-            nhwc_node = self._g.make_node("Transpose", [nchw_node.output[0]], attr={"perm": [0, 2, 3, 1]})
+            nchw_node = self._g.make_node("Transpose", [node.output[0]], attr={"perm": NHWC_TO_NCHW})
+            nhwc_node = self._g.make_node("Transpose", [nchw_node.output[0]], attr={"perm": NCHW_TO_NHWC})
             self._g.replace_input(consumer, node.output[0], nhwc_node.output[0])
 
     def _create_transpose_pairs_before_node(self, node):
@@ -425,7 +426,10 @@ def _identity_handler(self, trans, node):
 
     def _concat_handler(self, trans, node):
         if self._handle_node_having_branches(node):
-            node.set_attr("axis", 1)
+            perm = trans.get_attr_value("perm")
+            axis = node.get_attr_value("axis", 0)
+            new_axis = perm[axis]
+            node.set_attr("axis", new_axis)
             return True
         return False
 
@@ -505,7 +509,7 @@ def _slice_handler(self, trans, node):
                     axes = axes_node.get_tensor_value(as_list=True)
 
         if axes == [0, 1, 2, 3]:
-            node.set_attr("axes", [0, 2, 3, 1])
+            node.set_attr("axes", NCHW_TO_NHWC)
             return self._switch_transpose_and_node(node, trans)
         return False