fix issues in test_internals.py and onnx-experiments.py

Author: pengwa

tests/test_internals.py

@@ -19,7 +19,7 @@
 import tf2onnx
 import tf2onnx.utils
 from tf2onnx.graph_matcher import OpTypePattern, GraphMatcher
-from tf2onnx.graph import Graph
+from tf2onnx.graph import Graph, GraphUtil
 
 # pylint: disable=missing-docstring
 
@@ -43,8 +43,8 @@ def onnx_to_graphviz(g):
 
 def onnx_pretty(g, args=None):
     """Onnx graph pretty print."""
-    model_proto = g.make_model("converted from {}".format(args.input))
-    return helper.printable_graph(model_proto.graph)
+    graph_proto = g.make_model("converted from {}".format(args.input))
+    return helper.printable_graph(graph_proto.graph)
 
 
 class Tf2OnnxInternalTests(unittest.TestCase):
@@ -73,60 +73,63 @@ def sample_net():
         n5 = helper.make_node("Abs", ["n4:0"], ["n5:0"], name="n5")
         n6 = helper.make_node("Identity", ["n5:0"], ["n6:0"], name="n6")
 
-        model_proto = helper.make_graph(
+        graph_proto = helper.make_graph(
             nodes=[n1, n2, n3, n4, n5, n6],
             name="test",
             inputs=[helper.make_tensor_value_info("input", TensorProto.FLOAT, [2, 2])],
             outputs=[helper.make_tensor_value_info("n5:0", TensorProto.FLOAT, [2, 2])],
             initializer=[]
         )
-        return model_proto
+        return graph_proto
 
     def test_insert_node1(self):
-        model_proto = self.sample_net()
-        nodes = model_proto.node
-        g = Graph(nodes, output_shapes={}, dtypes={})
+        graph_proto = self.sample_net()
+        g = GraphUtil.create_graph_from_onnx_graph(graph_proto)
         n2 = g.get_node_by_name("n2")
         n7 = g.insert_new_node_on_input(n2, "Abs", "n1:0", name="n7")
         ops = g.get_nodes()
         ops.append(n7)
         g.topological_sort(ops)
         result = onnx_to_graphviz(g)
-        expected = 'digraph { n1 [op_type=Abs] n7 [op_type=Abs] n2 [op_type=Abs] n3 [op_type=Abs] ' \
-                   'n4 [op_type=Add] n5 [op_type=Abs] n6 [op_type=Identity] ' \
-                   'input -> n1 n1:0 -> n7 n7:0 -> n2 n1:0 -> n3 n2:0 -> n4 n3:0 -> n4 n4:0 -> n5 n5:0 -> n6 }'
+        expected = 'digraph { Placeholder__4 [op_type=Placeholder] ' \
+                   'n1 [op_type=Abs] n7 [op_type=Abs] n2 [op_type=Abs] n3 [op_type=Abs] ' \
+                   'n4 [op_type=Add] n5 [op_type=Abs] graph_outputs_Identity__3 [op_type=Identity] ' \
+                   'n6 [op_type=Identity] input -> n1 n1:0 -> n7 n7:0 -> n2 n1:0 -> n3 ' \
+                   'n2:0 -> n4 n3:0 -> n4 n4:0 -> n5 raw_output___2:0 -> graph_outputs_Identity__3 ' \
+                   'raw_output___2:0 -> n6 }'
         self.assertEqual(expected, result)
 
     def test_insert_node2(self):
-        model_proto = self.sample_net()
-        nodes = model_proto.node
-        g = Graph(nodes, output_shapes={}, dtypes={})
+        graph_proto = self.sample_net()
+        g = GraphUtil.create_graph_from_onnx_graph(graph_proto)
         n7 = g.insert_new_node_on_output("Abs", "n1:0", name="n7")
         ops = g.get_nodes()
         ops.append(n7)
         g.topological_sort(ops)
         result = onnx_to_graphviz(g)
-        expected = 'digraph { n1 [op_type=Abs] n7 [op_type=Abs] n3 [op_type=Abs] n2 [op_type=Abs] ' \
-                   'n4 [op_type=Add] n5 [op_type=Abs] n6 [op_type=Identity] ' \
-                   'input -> n1 n1:0 -> n7 n7:0 -> n3 n7:0 -> n2 n2:0 -> n4 n3:0 -> n4 n4:0 -> n5 n5:0 -> n6 }'
+        expected = 'digraph { Placeholder__4 [op_type=Placeholder] n1 [op_type=Abs] n7 [op_type=Abs] ' \
+                   'n3 [op_type=Abs] n2 [op_type=Abs] n4 [op_type=Add] n5 [op_type=Abs] ' \
+                   'graph_outputs_Identity__3 [op_type=Identity] n6 [op_type=Identity] ' \
+                   'input -> n1 n1:0 -> n7 n7:0 -> n3 n7:0 -> n2 n2:0 -> n4 n3:0 -> n4 ' \
+                   'n4:0 -> n5 raw_output___2:0 -> graph_outputs_Identity__3 raw_output___2:0 -> n6 }'
         self.assertEqual(expected, result)
 
     def test_remove_input(self):
-        model_proto = self.sample_net()
-        nodes = model_proto.node
-        g = Graph(nodes, output_shapes={}, dtypes={})
+        graph_proto = self.sample_net()
+        g = GraphUtil.create_graph_from_onnx_graph(graph_proto)
         n4 = g.get_node_by_name("n4")
         g.remove_input(n4, n4.input[1])
         result = onnx_to_graphviz(g)
         expected = 'digraph { n1 [op_type=Abs] n2 [op_type=Abs] n3 [op_type=Abs] n4 [op_type=Add] ' \
-                   'n5 [op_type=Abs] n6 [op_type=Identity] input -> n1 n1:0 -> n2 n1:0 -> n3 n2:0 -> n4 ' \
-                   'n4:0 -> n5 n5:0 -> n6 }'
+                   'n5 [op_type=Abs] n6 [op_type=Identity] graph_outputs_Identity__3 ' \
+                   '[op_type=Identity] Placeholder__4 [op_type=Placeholder] input -> n1 ' \
+                   'n1:0 -> n2 n1:0 -> n3 n2:0 -> n4 n4:0 -> n5 raw_output___2:0 -> n6 ' \
+                   'raw_output___2:0 -> graph_outputs_Identity__3 }'
         self.assertEqual(expected, result)
 
     def test_rewrite_subgraph(self):
-        model_proto = self.sample_net()
-        nodes = model_proto.node
-        g = tf2onnx.graph.Graph(nodes, output_shapes={}, dtypes={})
+        graph_proto = self.sample_net()
+        g = GraphUtil.create_graph_from_onnx_graph(graph_proto)
         pattern = \
             OpTypePattern('Abs', name='output', inputs=[
                 OpTypePattern('Add', name='input')
@@ -143,26 +146,28 @@ def test_rewrite_subgraph(self):
             ops = g.replace_subgraph(ops, match, [], [output_node], [], [new_node])
         g.topological_sort(ops)
         result = onnx_to_graphviz(g)
-        expected = 'digraph { n1 [op_type=Abs] n3 [op_type=Abs] n2 [op_type=Abs] ReplacedOp__2 [op_type=Sub] ' \
-                   'n6 [op_type=Identity] input -> n1 n1:0 -> n3 n1:0 -> n2 n2:0 -> ReplacedOp__2 ' \
-                   'n3:0 -> ReplacedOp__2 ReplacedOp__2:0 -> n6 }'
+        expected = 'digraph { Placeholder__4 [op_type=Placeholder] n1 [op_type=Abs] ' \
+                   'n3 [op_type=Abs] n2 [op_type=Abs] ReplacedOp__5 [op_type=Sub] ' \
+                   'graph_outputs_Identity__3 [op_type=Identity] n6 [op_type=Identity] ' \
+                   'input -> n1 n1:0 -> n3 n1:0 -> n2 n2:0 -> ReplacedOp__5 ' \
+                   'n3:0 -> ReplacedOp__5 ReplacedOp__5:0 -> graph_outputs_Identity__3 ' \
+                   'ReplacedOp__5:0 -> n6 }'
         self.assertEqual(expected, result)
 
     def test_match_flipped(self):
         n1 = helper.make_node("Sub", ["i1", "i1"], ["n1:0"], name="n1")
         n2 = helper.make_node("Add", ["i2", "i2"], ["n2:0"], name="n2")
         n3 = helper.make_node("Mul", ["n1:0", "n2:0"], ["n3:0"], name="n3")
 
-        model_proto = helper.make_graph(
+        graph_proto = helper.make_graph(
             nodes=[n1, n2, n3],
             name="test",
             inputs=[helper.make_tensor_value_info("i1", TensorProto.FLOAT, [2, 2]),
                     helper.make_tensor_value_info("i2", TensorProto.FLOAT, [2, 2])],
             outputs=[helper.make_tensor_value_info("n2:0", TensorProto.FLOAT, [2, 2])],
             initializer=[]
         )
-        nodes = model_proto.node
-        g = tf2onnx.graph.Graph(nodes, output_shapes={}, dtypes={})
+        g = GraphUtil.create_graph_from_onnx_graph(graph_proto)
         pattern = OpTypePattern('Mul', inputs=[
             OpTypePattern('Add'),
             OpTypePattern('Sub')

tf2onnx/graph.py

@@ -1041,6 +1041,10 @@ def create_graph_from_onnx_graph(graph_proto):
         for n in graph_proto.node:
             if n.op_type == "Constant":
                 n.op_type = "Const"
+
+            # some pytorch model had empty names - make one up
+            if not n.name:
+                n.name = utils.make_name("was_empty")
             nodes_to_append.append(n)
 
         output_names = []

tools/onnx-experiments.py

@@ -22,7 +22,7 @@
 from onnx import numpy_helper
 
 import tf2onnx.utils
-from tf2onnx.graph import Graph
+from tf2onnx.graph import Graph, GraphUtil
 
 logging.basicConfig(level=logging.INFO)
 log = logging.getLogger("onnx-experiments")
@@ -41,48 +41,7 @@ def load_graph(fname):
     with open(fname, "rb") as f:
         data = f.read()
         model_proto = onnx.ModelProto()
-        model_proto.ParseFromString(data)
-        onnx_nodes = model_proto.graph.node
-        output_names = []
-
-        # some pytorch model had empty names - make one up
-        for node in onnx_nodes:
-            if not node.name:
-                node.name = tf2onnx.utils.make_name("was_empty")
-
-        g = Graph(onnx_nodes, output_shapes={}, dtypes={}, output_names=output_names)
-        for i in model_proto.graph.initializer:
-            v = numpy_helper.to_array(i)
-            name = i.name
-            g.initializers[name] = i
-            dtype = i.data_type
-            g.set_dtype(name, dtype)
-            g.set_shape(name, v.shape)
-        for i in model_proto.graph.input:
-            name = i.name
-            if name in g.initializers:
-                # ignore if it is not a model input
-                continue
-            shape = [j.dim_value if hasattr(i.type.tensor_type, "dim_value") else -1
-                     for j in i.type.tensor_type.shape.dim]
-            dtype = i.type.tensor_type.elem_type
-            g.set_dtype(name, dtype)
-            g.set_shape(name, shape)
-            g.add_graph_input(name, dtype, shape)
-        for i in model_proto.graph.output:
-            name = i.name
-            shape = [j.dim_value if hasattr(i.type.tensor_type, "dim_value") else -1
-                     for j in i.type.tensor_type.shape.dim]
-            dtype = i.type.tensor_type.elem_type
-            g.set_dtype(name, dtype)
-            g.set_shape(name, shape)
-            output_names.append(name)
-
-        # TODO: this is a hack in case a output name does not follow tensorflow convention
-        for node in g.get_nodes():
-            for name in node.output:
-                g._nodes_by_name[name] = node  # pylint: disable=protected-access
-
+        g = GraphUtil.create_graph_from_onnx_model(model_proto)
     return g, model_proto.producer_name