refactor

Author: zhijxu-MS

tf2onnx/optimizer/transpose_optimizer.py

@@ -7,9 +7,8 @@
 from collections import defaultdict
 
 import numpy as np
-
-from tf2onnx.constants import NCHW_TO_NHWC, NHWC_TO_NCHW
 import onnx
+from tf2onnx.constants import NCHW_TO_NHWC, NHWC_TO_NCHW
 from .. import utils
 from .optimizer_base import GraphOptimizerBase
 
@@ -192,8 +191,8 @@ def _initialize_handlers(self):
 
     def _handle_node_having_branches(self, node):
         # create transpose pairs if some input are not.
-        self._create_transpose_pairs_before_node(node)
-
+        if not self._create_transpose_pairs_before_node(node):
+            return False
         # make sure node's all input transpose all have only 1 consumer node,
         # otherwise, it would impact their other output nodes
         if self._nodes_has_single_consumer_node(node.inputs):
@@ -311,9 +310,10 @@ def _create_transpose_pairs_before_node(self, node):
         def shape_after_expand(ori_shape):
             # according to broadcasting rule to expand shape to 4D while not tile the tensor here
             # still count on the broadcasting op to tile the tensor
-            utils.make_sure(ori_shape.count(-1) <= 1, "shape can contain one -1 at most")
+            if ori_shape.count(-1) >= 2:
+                self.logger.warning("%s shape can contain one -1 at most, otherwise reshape op can't work", node.name)
+                return None
             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
 
@@ -325,6 +325,14 @@ def shape_after_expand(ori_shape):
                     non_nhwc_trans_inputs.append([input_id, n])
 
         # add Transpose(0, 3, 1, 2) and Transpose(0, 2, 3, 1) before each non_nhwc_trans_consumers
+        shape_unknow = [input_id for input_id, _ in non_nhwc_trans_inputs if self._g.get_shape(input_id) is None]
+        if shape_unknow:
+            if self._g.opset <= 9:
+                msg = "%s 's shape is unknown, ConstantOfShape will be used which exists in version 9 or higher" \
+                      "while graph's opset version is %s" % (shape_unknow, self._g.opset)
+                self.logger.warning(msg)
+                return False
+
         for input_id, n in non_nhwc_trans_inputs:
             shape = self._g.get_shape(input_id)
             # if rank of n is not 4, then we need to insert a reshape op before inserting a transpose
@@ -343,13 +351,16 @@ def shape_after_expand(ori_shape):
                 input_of_new_trans = input_id
             else:
                 shape_4d = shape_after_expand(shape)
+                if shape_4d is None:
+                    return False
                 const = self._g.make_const(utils.make_name("reshape_shape"), np.array(shape_4d, np.int64)).output[0]
                 reshape = self._g.make_node("Reshape", [input_id, const]).output[0]
                 input_of_new_trans = reshape
 
-            nchw_node = self._g.make_node("Transpose", [input_of_new_trans], 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", [input_of_new_trans], 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(node, input_id, nhwc_node.output[0])
+        return True
 
     def _add_handler(self, trans, node):
         if node.inputs[1].is_const():