Hash tensor tensor values in merge_duplicated_nodes to increase conversion speed. Decreases total time from 16 min to 6 min for bit/m-r152x4. For small models time is unchanged

Author: TomWildenhain-Microsoft

tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

@@ -7,16 +7,14 @@
    then b and c can be merged into one node to avoid duplicated computation
 """
 
-from collections import defaultdict, namedtuple
+from collections import defaultdict
 
 import numpy as np
 
 from .optimizer_base import GraphOptimizerBase
 
 # pylint: disable=logging-not-lazy,unused-argument,missing-docstring
 
-_KeyToGroupNodes = namedtuple("key", "type input")
-
 
 class MergeDuplicatedNodesOptimizer(GraphOptimizerBase):
     """Remove duplicate nodes.
@@ -41,6 +39,7 @@ def _optimize_at_current_graph_level(self, graph):
     def _merge_duplicated_nodes(self, graph):
         # "duplicated" means: op_type, input and attribute are same
         # while attr is un-hashable so doesn't include it when grouping nodes
+        # we do hash the tensor data of const values
         nodes_groups = self._group_nodes_by_type_inputs(graph)
         for _, nodes_group in nodes_groups.items():
             if self._skip_node_type(nodes_group[0]):
@@ -54,7 +53,11 @@ def _group_nodes_by_type_inputs(graph):
             # default const of graph input cannot be merged
             if node.is_graph_input_default_const():
                 continue
-            res[_KeyToGroupNodes(node.type, tuple(node.input))].append(node)
+            tensor_data_hash = None
+            if node.is_const():
+                # Many constants have the same size so this is helpful
+                tensor_data_hash = hash(node.attr['value'].t.raw_data)
+            res[(node.type, tuple(node.input), tensor_data_hash)].append(node)
         return res
 
     def _del_nodes_if_duplicated(self, nodes_group, graph):
@@ -75,7 +78,7 @@ def _del_nodes_if_duplicated(self, nodes_group, graph):
     def _have_equal_attr(self, node_1, node_2, graph):
         if node_1.attr == node_2.attr:
             return True
-        # above check guarantees consts here are able to be merged
+        # consts have a name attr that can differ among equal consts so they must be handled separately
         if node_1.is_const() and node_2.is_const():
             # get_tensor_value is costly so that we check their shape first
             shape_1 = graph.get_shape(node_1.output[0])