refactor

use NHWC_TO_NCHW instead of [0, 3, 1, 2]

Author: zhijxu-MS

tf2onnx/onnx_opset/nn.py

@@ -538,7 +538,7 @@ 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

tf2onnx/optimizer/transpose_optimizer.py

@@ -24,7 +24,7 @@ def is_nhwc_transpose(transpose_node):
 
 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):
@@ -87,7 +87,7 @@ 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(NCHW_TO_NHWC)).output[0]
 
@@ -246,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
 
@@ -302,7 +302,7 @@ 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]})
+            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])