Merge pull request #1102 from onnx/tom/ImproveOptDupNodePerf

Hash tensor tensor values in merge_duplicated_nodes to increase conversion speed

Author: TomWildenhain-Microsoft

tests/test_optimizers.py

@@ -918,8 +918,8 @@ def test_duplicated_duplicated_constant(self):
         const_val = np.array([1, 2, 3], dtype=np.float32)
         tensor_1 = helper.make_tensor("tensor_1", TensorProto.FLOAT, const_val.shape, const_val)
         tensor_2 = helper.make_tensor("tensor_2", TensorProto.FLOAT, const_val.shape, const_val)
-        tensor_3 = helper.make_tensor("tensor_3", TensorProto.FLOAT, const_val.shape, const_val.tobytes(), raw=True)
-        tensor_4 = helper.make_tensor("tensor_4", TensorProto.FLOAT, const_val.shape, const_val.tobytes(), raw=True)
+        tensor_3 = helper.make_tensor("tensor_3", TensorProto.FLOAT, const_val.shape, const_val)
+        tensor_4 = helper.make_tensor("tensor_4", TensorProto.FLOAT, const_val.shape, const_val)
         node0 = helper.make_node('Constant', inputs=[], outputs=["value0"], value=tensor_1)
         node1 = helper.make_node('Constant', inputs=[], outputs=["value1"], value=tensor_2)
         node2 = helper.make_node('Constant', inputs=[], outputs=["value2"], value=tensor_3)
@@ -941,8 +941,8 @@ def test_duplicated_duplicated_constant(self):
 
     def test_duplicated_duplicated_constant_and_initializer(self):
         const_val = np.array([1, 2, 3], dtype=np.float32)
-        tensor_1 = helper.make_tensor("value0", TensorProto.FLOAT, const_val.shape, const_val)
-        tensor_2 = helper.make_tensor("value1", TensorProto.FLOAT, const_val.shape, const_val)
+        tensor_1 = helper.make_tensor("value0", TensorProto.FLOAT, const_val.shape, const_val.tobytes(), raw=True)
+        tensor_2 = helper.make_tensor("value1", TensorProto.FLOAT, const_val.shape, const_val.tobytes(), raw=True)
         tensor_3 = helper.make_tensor("value2", TensorProto.FLOAT, const_val.shape, const_val.tobytes(), raw=True)
         tensor_4 = helper.make_tensor("value3", TensorProto.FLOAT, const_val.shape, const_val.tobytes(), raw=True)
         node0 = helper.make_node('Constant', inputs=[], outputs=["value0"], value=tensor_1)

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])