Merge pull request #744 from jignparm/jignparm/bert_model_cleanup

Remove unnecessary ops when converting BERT model

Author: jignparm

tf2onnx/onnx_opset/tensor.py

@@ -722,23 +722,24 @@ def version_1(cls, ctx, node, **kwargs):
 
         # onnx slice as of opset 7 does only take float tensors ... cast if needed
         input_dtype = ctx.get_dtype(node.input[0])
-        if input_dtype != onnx_pb.TensorProto.FLOAT:
-            if node.inputs[0].type == "Cast" and len(ctx.find_output_consumers(node.inputs[0].output[0])) == 1:
-                # override the previous cast
-                cast_node = node.inputs[0]
-            else:
-                cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
-                nodes.insert(0, cast_node)
-            cast_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
-            ctx.set_dtype(cast_node.output[0], onnx_pb.TensorProto.FLOAT)
-            ctx.copy_shape(node.input[0], cast_node.output[0])
-            # undo the cast afer slice
-            name = utils.make_name(node.name)
-            cast_node = ctx.insert_new_node_on_output("Cast", nodes[-1].output[0], name)
-            cast_node.set_attr("to", input_dtype)
-            ctx.set_dtype(cast_node.output[0], input_dtype)
-            ctx.copy_shape(node.output[0], cast_node.output[0])
-            nodes.append(cast_node)
+        if ctx.opset < 9:
+            if input_dtype != onnx_pb.TensorProto.FLOAT:
+                if node.inputs[0].type == "Cast" and len(ctx.find_output_consumers(node.inputs[0].output[0])) == 1:
+                    # override the previous cast
+                    cast_node = node.inputs[0]
+                else:
+                    cast_node = ctx.insert_new_node_on_input(node, "Cast", node.input[0])
+                    nodes.insert(0, cast_node)
+                cast_node.set_attr("to", onnx_pb.TensorProto.FLOAT)
+                ctx.set_dtype(cast_node.output[0], onnx_pb.TensorProto.FLOAT)
+                ctx.copy_shape(node.input[0], cast_node.output[0])
+                # undo the cast afer slice
+                name = utils.make_name(node.name)
+                cast_node = ctx.insert_new_node_on_output("Cast", nodes[-1].output[0], name)
+                cast_node.set_attr("to", input_dtype)
+                ctx.set_dtype(cast_node.output[0], input_dtype)
+                ctx.copy_shape(node.output[0], cast_node.output[0])
+                nodes.append(cast_node)
 
     @classmethod
     def version_10(cls, ctx, node, **kwargs):

tf2onnx/optimizer/back_to_back_optimizer.py

@@ -36,7 +36,7 @@ def _optimize(self, graph):
     def _optimize_at_current_graph_level(self, g):
         for optype, handler in _func_map.items():
             # candidate nodes for removal/optimization
-            nodes = [n for n in g.get_nodes() if n.type == optype]
+            nodes = [n for n in g.get_nodes() if n.type in optype]
 
             # topological sort of candidates
             # simplifying assumption for back-to-back-optimizer is
@@ -51,8 +51,7 @@ def _optimize_at_current_graph_level(self, g):
             # q = starting nodes with no dependencies
             q = list(set(consumer_node_ids.keys()) - has_dependencies)
             while q:
-                nodeid = q[0]
-                q.remove(nodeid)
+                nodeid = q.pop(0)
                 node = g.get_node_by_output(nodeid, False)
                 consumer_nodes = consumer_node_ids[nodeid]
 
@@ -72,6 +71,7 @@ def _optimize_at_current_graph_level(self, g):
     @staticmethod
     @_register_func("Cast")
     def _optimize_cast(g, node, consumer_nodes):
+        """remove long chains of cast ops"""
         q2 = []
         type1 = node.get_attr('to').i
         type1_name = ONNX_DTYPE_NAMES[type1] if type1 in ONNX_DTYPE_NAMES else ''
@@ -124,6 +124,7 @@ def _optimize_cast(g, node, consumer_nodes):
     @staticmethod
     @_register_func("Transpose")
     def _optimize_transpose(g, node, consumer_nodes):
+        """remove long chains of transpose ops"""
         t1 = list(node.get_attr('perm').ints)
         q2 = []
         for node2 in consumer_nodes:
@@ -146,3 +147,33 @@ def _optimize_transpose(g, node, consumer_nodes):
                 q2.append(node2.output[0])
         g.remove_node(node.name)
         return q2
+
+    @staticmethod
+    @_register_func(('Squeeze', 'Unsqueeze'))
+    def _optimize_squeeze_unsqueeze(g, node, consumer_nodes):
+        """remove pairs of squeeze-unsqueeze nodes"""
+
+        if node.type != 'Squeeze' or len(consumer_nodes) != 1:
+            # no need to return any value, since not removing long chain of nodes
+            return []
+
+        node2 = consumer_nodes[0]
+        if node2.type != 'Unsqueeze':
+            return []
+
+        axis1 = node.get_attr('axes').ints
+        axis2 = node2.get_attr('axes').ints
+
+        # if squeeze followed by unsqueeze is on diff axes, skip
+        if axis1 != axis2:
+            return []
+
+        # if unsqueeze output is graph output, skip
+        if set(node2.output) & set(g.outputs):
+            return []
+
+        node2_consumers = g.find_output_consumers(node2.output[0])
+        g.replace_all_inputs(node2_consumers, node2.output[0], node.input[0])
+        g.remove_node(node.name)
+        g.remove_node(node2.name)
+        return []