merge const

Author: wayuanho

tests/test_optimizers.py

@@ -12,7 +12,7 @@
 from tf2onnx import utils
 from tf2onnx.graph import GraphUtil
 from backend_test_base import Tf2OnnxBackendTestBase
-from common import unittest_main, group_nodes_by_type, check_opset_min_version
+from common import unittest_main, group_nodes_by_type
 
 
 # pylint: disable=missing-docstring,invalid-name,unused-argument,using-constant-test
@@ -758,9 +758,12 @@ def test_identity_in_subgraph_non_graph_output(self):
     # Merge Duplicated Nodes Optimizer Tests Start
 
     def run_merge_duplicated_nodes_compare(self, output_names_with_port, onnx_feed_dict, origin_proto,
-                                           op_type=None, remaining_op_num=None, debug=False, rtol=1e-07):
-        self.run_and_compare(output_names_with_port, onnx_feed_dict, origin_proto, op_type=op_type,
-                             remaining_op_num=remaining_op_num, debug=debug, rtol=rtol)
+                                           op_type=None, remaining_op_num=None, debug=False, rtol=1e-07,
+                                           graph_validator=None):
+        new_proto = self.run_and_compare(output_names_with_port, onnx_feed_dict, origin_proto, op_type=op_type,
+                                         remaining_op_num=remaining_op_num, debug=debug, rtol=rtol)
+        if graph_validator:
+            self.assertTrue(graph_validator(new_proto.graph))
 
     def test_duplicated_duplicated_input(self):
         # same input or not
@@ -800,7 +803,10 @@ def test_duplicated_duplicated_attributes(self):
         self.run_merge_duplicated_nodes_compare(["OUT"], {"X": np.random.randn(5, 5).astype(np.float32)}, model_proto,
                                                 op_type="ReduceSum", remaining_op_num=2)
 
-    @check_opset_min_version(9, "Constant")
+    def _check_initializer_num(self, graph_proto, num):
+        print(len(graph_proto.initializer))
+        return num == len(graph_proto.initializer)
+
     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)
@@ -826,8 +832,35 @@ def test_duplicated_duplicated_constant(self):
         imp.version = self.config.opset
 
         model_proto = helper.make_model(graph, producer_name="onnx-tests", opset_imports=[imp])
-        self.run_merge_duplicated_nodes_compare(["OUT"], {}, model_proto,
-                                                op_type="Constant", remaining_op_num=1)
+        self.run_merge_duplicated_nodes_compare(["OUT"], {}, model_proto, op_type="Constant", remaining_op_num=0,
+                                                graph_validator=lambda g: self._check_initializer_num(g, 1))
+
+    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_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)
+        node1 = helper.make_node('Constant', inputs=[], outputs=["value1"], value=tensor_2)
+        node4 = helper.make_node("Mul", ["value0", "value1"], ["output1"])
+        node5 = helper.make_node("Mul", ["value2", "output1"], ["output2"])
+        node6 = helper.make_node("Mul", ["value3", "output2"], ["OUT"])
+
+        graph = helper.make_graph(
+            [node0, node1, node4, node5, node6],
+            "test_duplicated_duplicated_constant",
+            [helper.make_tensor_value_info("value2", TensorProto.FLOAT, (3,))],
+            [helper.make_tensor_value_info("OUT", TensorProto.FLOAT, (3,))],
+            [tensor_3, tensor_4]
+        )
+
+        imp = OperatorSetIdProto()
+        imp.version = self.config.opset
+
+        model_proto = helper.make_model(graph, producer_name="onnx-tests", opset_imports=[imp])
+        self.run_merge_duplicated_nodes_compare(["OUT"], {}, model_proto, op_type="Constant", remaining_op_num=0,
+                                                graph_validator=lambda g: self._check_initializer_num(g, 2))
 
     def test_duplicated_node_is_graph_output(self):
         node0 = helper.make_node('Add', inputs=["X", "X"], outputs=["value0"])

tf2onnx/graph.py

@@ -17,7 +17,7 @@
 
 from onnx import helper, numpy_helper, shape_inference, OperatorSetIdProto, AttributeProto, TensorProto
 from tf2onnx import utils, __version__
-from tf2onnx.utils import port_name, find_opset
+from tf2onnx.utils import make_name, port_name, find_opset
 from tf2onnx import optimizer
 from tf2onnx.schemas import get_schema, infer_onnx_shape_dtype
 from tf2onnx import constants
@@ -146,6 +146,11 @@ def is_const(self):
     def is_graph_input(self):
         return self.type in ["Placeholder", "PlaceholderWithDefault", "PlaceholderV2"]
 
+    def is_graph_input_default_const(self):
+        return self.is_const() and any(
+            out.is_graph_input() for out in self.graph.find_output_consumers(self.output[0])
+        )
+
     def __str__(self):
         return str(self._op)
 
@@ -711,6 +716,20 @@ def add_graph_input(self, name, dtype=None, shape=None):
         new_node = self.make_node("Placeholder", [], outputs=[name], dtypes=[dtype], shapes=[shape])
         self._order_sensitive_inputs.append(new_node)
 
+    def add_graph_input_with_default(self, name, default_const, dtype=None, shape=None):
+        """Add placeholderwithdefault."""
+        if dtype is None:
+            dtype = self.get_dtype(name)
+
+        if shape is None:
+            shape = self.get_shape(name)
+
+        default_const_name = port_name(make_name("{}_default".format(name)))
+        default_const.output = [default_const_name]
+        new_node = self.make_node("PlaceholderWithDefault", [default_const_name], outputs=[name],
+                                  dtypes=[dtype], shapes=[shape])
+        self._order_sensitive_inputs.append(new_node)
+
     def add_graph_output(self, name, dtype=None, shape=None):
         """Add node output as graph's output."""
         utils.make_sure(name in self._output_to_node_name, "output %s not exist in the graph", name)
@@ -765,6 +784,8 @@ def set_shape(self, name, val):
         """Set new shape of node."""
         if isinstance(val, np.ndarray):
             val = val.tolist()
+        if isinstance(val, tuple):
+            val = list(val)
         node = self.get_node_by_output(name, search_in_parent_graphs=True)
         utils.make_sure(node is not None, "cannot find node by output id %s", name)
         node.graph._output_shapes[name] = val
@@ -1358,3 +1379,5 @@ def _parse_graph_input(g, graph_proto, const_node_names):
             dtype = dtypes[name]
             if name not in const_node_names:
                 g.add_graph_input(name, dtype, shape)
+            else:
+                g.add_graph_input_with_default(name, g.get_node_by_name(name), dtype, shape)

tf2onnx/optimizer/merge_duplicated_nodes_optimizer.py

@@ -51,6 +51,9 @@ def _merge_duplicated_nodes(self, graph):
     def _group_nodes_by_type_inputs(graph):
         res = defaultdict(list)
         for node in graph.get_nodes():
+            # 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)
         return res
 
@@ -72,6 +75,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
         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])