merge transpose nodes which share same input and perm attr are also same

current strategy is to move down transpose nodes if possible while this is
is not always a good idea for it may generate more transpose nodes which
not removed in later procedures. so it's better to merge them back into
one after removing procedures.

Author: zhijxu-MS

tf2onnx/optimizer/transpose_optimizer.py

@@ -3,6 +3,7 @@
 """Transpose Optimizer."""
 
 from __future__ import unicode_literals
+from collections import defaultdict
 
 import logging
 
@@ -113,6 +114,27 @@ def post_optimize_action(self):
         self._g.update_proto()
         self._g.topological_sort(self._g.get_nodes())
 
+    def merge_transpose_with_same_input(self):
+        # strategy used in previous procedure is to move transpose nodes down if possible,
+        # and it means that when a node has n outputs then n transpose will be generated ,
+        # so we should merge them back to one if they can't be eliminated in previous procedure.
+        graph = self._g
+        input_transposes_map = defaultdict(list)
+        for node in graph.get_nodes():
+            if node.type == "Transpose":
+                key = (node.input[0], str(node.get_attr("perm").ints))
+                input_transposes_map[key].append(node)
+
+        for _, transposes in input_transposes_map.items():
+            # merge transpose nodes into one: make nodes use the output of the first transpose node
+            transpose_out = transposes[0].output[0]
+            for node in transposes[1:]:
+                old_transpose_out = node.output[0]
+                graph.replace_all_inputs(graph.get_nodes(), old_transpose_out, transpose_out)
+
+        # dangling transpose nodes can be deleted
+        graph.delete_unused_nodes(graph.outputs)
+
     def optimize(self):
         previous_counter = self._g.dump_node_statistics()
         no_action = False
@@ -140,7 +162,9 @@ def optimize(self):
                 break
 
         log.debug("finish after " + str(iteration_cnt) + " iteration(s)")
+
         self.post_optimize_action()
+        self.merge_transpose_with_same_input()
 
         current_counter = self._g.dump_node_statistics()
         transpose_cnt = current_counter["Transpose"]