Merge pull request #340 from pengwa/set_nodes_change

Graph node management in real time when add/delete

Author: pengwa

tests/run_pretrained_models.py

@@ -265,6 +265,7 @@ def run_test(self, name, backend="caffe2", debug=False, onnx_file=None, opset=No
                 # convert model to onnx
                 onnx_graph = self.to_onnx(sess.graph, opset=opset, shape_override=shape_override,
                                           input_names=inputs.keys())
+                model_proto = onnx_graph.make_model("converted from tf2onnx")
                 new_model_proto = GraphUtil.opt_transposes_with_graph(onnx_graph, "test", debug=debug)
                 if new_model_proto:
                     model_proto = new_model_proto

tests/run_pretrained_models.yaml

@@ -91,7 +91,7 @@ inception_v3_slim:
   outputs:
     - InceptionV3/Predictions/Softmax:0
   rtol: 0.02
-  atol: 0.000001
+  atol: 0.00001
 
 googlenet_v1_nonslim:
   disabled: true

tests/test_backend.py

@@ -474,7 +474,7 @@ def test_matmul3(self):
         y = tf.placeholder(tf.float32, x_shape, name=_TFINPUT1)
         x_ = tf.matmul(x, y, transpose_b=True)
         _ = tf.identity(x_, name=_TFOUTPUT)
-        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: x_val}, rtol=1e-6)
+        self._run_test_case([_OUTPUT], {_INPUT: x_val, _INPUT1: x_val}, rtol=1e-5)
 
     @check_onnxruntime_incompatibility("Sub")
     def test_sub(self):

tests/test_internals.py

@@ -86,9 +86,8 @@ def test_insert_node1(self):
         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")
+        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 { Placeholder__4 [op_type=Placeholder] ' \
@@ -102,9 +101,8 @@ def test_insert_node1(self):
     def test_insert_node2(self):
         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")
+        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 { Placeholder__4 [op_type=Placeholder] n1 [op_type=Abs] n7 [op_type=Abs] ' \
@@ -119,12 +117,14 @@ def test_remove_input(self):
         g = GraphUtil.create_graph_from_onnx_graph(graph_proto)
         n4 = g.get_node_by_name("n4")
         g.remove_input(n4, n4.input[1])
+        ops = g.get_nodes()
+        g.topological_sort(ops)
         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] 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 }'
+        expected = 'digraph { Placeholder__4 [op_type=Placeholder] n1 [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 -> n3 n1:0 -> n2 n2:0 -> n4 n4:0 -> n5 ' \
+                   'raw_output___2:0 -> graph_outputs_Identity__3 raw_output___2:0 -> n6 }'
         self.assertEqual(expected, result)
 
     def test_rewrite_subgraph(self):
@@ -143,7 +143,9 @@ def test_rewrite_subgraph(self):
             op_name = utils.make_name("ReplacedOp")
             out_name = utils.port_name(op_name)
             new_node = g.make_node("Sub", inputs=input_node.input, outputs=[out_name], name=op_name)
-            ops = g.replace_subgraph(ops, match, [], [output_node], [], [new_node])
+            g.replace_all_inputs(ops, output_node.output[0], new_node.output[0])
+            for n in set(match.get_nodes()):
+                g.remove_node(n.name)
         g.topological_sort(ops)
         result = onnx_to_graphviz(g)
         expected = 'digraph { Placeholder__4 [op_type=Placeholder] n1 [op_type=Abs] ' \

tf2onnx/function/gathernd.py

@@ -19,11 +19,8 @@ def _make_gathernd_inner_loop(ctx, params, index, dtype):
     # for (int i = 0; i < size(index); i++)
     #   gather_res = gather(gather_cur, index[i])
     scope_name = utils.make_name("gathernd_inner_loop")
-    nodes = []
     trip_node = ctx.make_node("Size", [index.output[0]])
-    nodes.append(trip_node)
     cond_const = ctx.make_const(utils.make_name("cond"), np.ones((), dtype=np.bool))
-    nodes.append(cond_const)
     trip_name = utils.make_name("i")
     cond_name = utils.make_name("cond")
     cond_out_name = utils.make_name("cond_out")
@@ -34,9 +31,8 @@ def _make_gathernd_inner_loop(ctx, params, index, dtype):
     g = ctx.create_new_graph_with_same_config()
     index_i = g.make_node("Gather", [index.output[0], trip_name], attr={"axis": 0})
     gather = g.make_node("Gather", [cur_name, index_i.output[0]], attr={"axis": 0})
-    squeeze = g.make_node("Squeeze", [gather.output[0]], attr={"axes": [0]}, outputs=[result_name])
-    cond = g.make_node("Identity", [cond_name], outputs=[cond_out_name])
-    g.set_nodes([index_i, gather, squeeze, cond])
+    g.make_node("Squeeze", [gather.output[0]], attr={"axes": [0]}, outputs=[result_name])
+    g.make_node("Identity", [cond_name], outputs=[cond_out_name])
 
     g.add_graph_input(trip_name, TensorProto.INT64, [])
     g.add_graph_input(cond_name, TensorProto.BOOL, [])
@@ -50,15 +46,13 @@ def _make_gathernd_inner_loop(ctx, params, index, dtype):
                                         params],
                                op_name_scope=scope_name, skip_conversion=False)
     inner_loop.set_body_graph_as_attr("body", g)
-    nodes.append(inner_loop)
-    return nodes, inner_loop
+    return inner_loop
 
 
 def make_gathernd(ctx, params, indices, output, scope_name, t_params):
     """make GatherNd op."""
     # Tparams output = GatherNd(Tparams params, Tidx indices)
     scope_name = utils.make_name(scope_name)
-    nodes = []
     # reshape indices into [sum(indices[:-1]), indices[-1]]
     indices_shape = ctx.make_node("Shape", [indices], dtypes=[TensorProto.INT64])
     indices_size = ctx.make_node("Size", [indices])
@@ -74,12 +68,11 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params):
                                   attr={"axis": 0},
                                   dtypes=[TensorProto.INT64])
     flatten_indices = ctx.make_node("Reshape", [indices, flatten_shape.output[0]])
-    nodes.extend([indices_shape, indices_size, inner_shape, outter_shape, flatten_shape, flatten_indices])
+
     # outter loop for each index
     # for (int i=0; i<outter_shape; i++) inner_loop(params, flatten_indices[i])
     cond_const = ctx.make_const(utils.make_name("cond"), np.ones((), dtype=np.bool))
     dummy_const = ctx.make_const(utils.make_name("dummy"), np.ones((), dtype=np.int64))
-    nodes.extend([cond_const, dummy_const])
 
     # body graph creation
     g = ctx.create_new_graph_with_same_config()
@@ -93,12 +86,10 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params):
     index = g.make_node("Gather", [flatten_indices.output[0], trip_name], attr={"axis": 0})
     index_squeeze = g.make_node("Squeeze", [index.output[0]], attr={"axes": [0]})
     # inner loop to gather result
-    nodes_to_append, inner_loop = _make_gathernd_inner_loop(g, params, index_squeeze, t_params)
-    g.set_nodes(nodes_to_append +
-                [index, index_squeeze,
-                 g.make_node("Identity", [cond_name], outputs=[cond_out_name]),
-                 g.make_node("Identity", [dummy_name], outputs=[dummy_out_name]),
-                 g.make_node("Identity", [inner_loop.output[0]], outputs=[result_name])])
+    inner_loop = _make_gathernd_inner_loop(g, params, index_squeeze, t_params)
+    g.make_node("Identity", [cond_name], outputs=[cond_out_name])
+    g.make_node("Identity", [dummy_name], outputs=[dummy_out_name])
+    g.make_node("Identity", [inner_loop.output[0]], outputs=[result_name])
 
     g.add_graph_input(trip_name, TensorProto.INT64, [])
     g.add_graph_input(cond_name, TensorProto.BOOL, [])
@@ -113,7 +104,7 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params):
                                   output_count=2,
                                   op_name_scope=scope_name, skip_conversion=False)
     gathernd_loop.set_body_graph_as_attr("body", g)
-    nodes.append(gathernd_loop)
+
     # reshape to target shape
     # output shape of gathernd: indices.shape[:-1] + gathernd_output.shape[1:]
     inner_loop_shape = ctx.make_node("Shape", [gathernd_loop.output[1]], dtypes=[TensorProto.INT64])
@@ -140,18 +131,7 @@ def make_gathernd(ctx, params, indices, output, scope_name, t_params):
                                  [output_shape_.output[0]],
                                  attr={"axes": [0], "ends": [-1], "starts": [0]},
                                  dtypes=[TensorProto.INT64])
-    output_reshape = ctx.make_node("Reshape",
-                                   [gathernd_loop.output[1], output_shape.output[0]],
-                                   outputs=[output])
-    nodes.extend([indices_outter_shape,
-                  inner_loop_shape,
-                  one_const,
-                  inner_loop_shape_,
-                  output_inner_shape,
-                  output_shape_,
-                  output_shape,
-                  output_reshape])
-    return nodes
+    ctx.make_node("Reshape", [gathernd_loop.output[1], output_shape.output[0]], outputs=[output])
 
 
 def gathernd_op(ctx, node, name, args):
@@ -163,4 +143,4 @@ def gathernd_op(ctx, node, name, args):
     # same as the attr Tparams
     t_params = ctx.get_dtype(params)
     utils.make_sure(t_params, "Dtype of {} is None".format(indices))
-    return make_gathernd(ctx, params, indices, output, name, t_params)
+    make_gathernd(ctx, params, indices, output, name, t_params)

tf2onnx/function/matrixbandpart.py

@@ -20,29 +20,22 @@ def matrixbandpart_op(ctx, node, name, args):
     # methods to generate mask matrix: if lower triangular is needed, then generate column one by one
     # otherwise row is generated one by one.
     axis, counter_axis, squeeze_axis = (1, 0, 2) if bandpart == [-1, 0] else (0, 1, 1)
-    nodes = []
     # 1: subgraph to implement tf.onelike(input[:, 0]),
     # no need to worry about the dtype, because bool type is needed as Xor only support bool
     node_name = utils.make_name("const_zero")
     const_zero = ctx.make_const(name=node_name, np_val=np.array([0]).astype(np.int32))
-    nodes.append(const_zero)
     first_col_or_row = ctx.make_node(op_type="Gather", inputs=[node.input[0], const_zero.output[0]],
                                      attr={"axis": axis})
-    nodes.append(first_col_or_row)
     first_col_or_row_casted = ctx.make_node(op_type="Cast", inputs=first_col_or_row.output,
                                             attr={"to": onnx_pb.TensorProto.BOOL})
-    nodes.append(first_col_or_row_casted)
     # line means one col or one row
     zero_line = ctx.make_node(op_type="Xor", inputs=first_col_or_row_casted.output*2)
-    nodes.append(zero_line)
     one_line = ctx.make_node(op_type="Not", inputs=zero_line.output)
-    nodes.append(one_line)
 
     # 2: "loop" to generate mask matrix: generate col or row of matrix one by one
     g = ctx.create_new_graph_with_same_config()
     node_name = utils.make_name("const_zero_bool")
     const_zero_bool = ctx.make_const(name=node_name, np_val=np.array([[0]]).astype(np.bool))
-    nodes.append(const_zero_bool)
     ctx.set_dtype(const_zero_bool.output[0], onnx_pb.TensorProto.BOOL)
 
     # shift right the line and add zero at the left.
@@ -51,11 +44,10 @@ def matrixbandpart_op(ctx, node, name, args):
     slice_node = g.make_node(op_type="Slice", inputs=[new_line.output[0]],
                              attr={"axes": [counter_axis], "starts": [0], "ends": [-1]})
 
-    body_nodes = [slice_node, new_line,
-                  g.make_node("Identity", ["cond"], outputs=["cond_out"]),
-                  g.make_node("Identity", ["line"], outputs=["res"]),
-                  g.make_node("Identity", [slice_node.output[0]], outputs=["line_out"])]
-    g.set_nodes(body_nodes)
+    g.make_node("Identity", ["cond"], outputs=["cond_out"])
+    g.make_node("Identity", ["line"], outputs=["res"])
+    g.make_node("Identity", [slice_node.output[0]], outputs=["line_out"])
+
     g.add_graph_input("trip", onnx_pb.TensorProto.INT64, [])
     g.add_graph_input("cond", onnx_pb.TensorProto.BOOL, [])
     g.add_graph_input("line", onnx_pb.TensorProto.BOOL, [-1, -1])
@@ -66,35 +58,28 @@ def matrixbandpart_op(ctx, node, name, args):
 
     # initial value of body vars
     shape = ctx.make_node(op_type="Shape", inputs=[node.input[0]])  # dtype of result is int64
-    nodes.append(shape)
     node_name = utils.make_name("line_num_index")
     col_or_row_num_index = ctx.make_const(name=node_name, np_val=np.array(axis).astype(np.int32))
-    nodes.append(col_or_row_num_index)
     line_num = ctx.make_node(op_type="Gather", inputs=[shape.output[0], col_or_row_num_index.output[0]])
-    nodes.append(line_num)
     trip_cnt = line_num.output[0]
     node_name = utils.make_name("true")
     cond = ctx.make_const(name=node_name, np_val=np.array(1).astype(np.bool))
-    nodes.append(cond)
     col_init = one_line.output[0]
 
     loop_node = ctx.make_node(op_type="Loop", inputs=[trip_cnt, cond.output[0], col_init], output_count=2)
     loop_node.set_body_graph_as_attr("body", g)
-    nodes.append(loop_node)
     # convert generated mask matrix from bool to right shape and data type
     squeeze = ctx.make_node(op_type="Squeeze", inputs=[loop_node.output[1]], attr={"axes": [squeeze_axis]})
-    nodes.append(squeeze)
     cast1 = ctx.make_node(op_type="Cast", inputs=squeeze.output, attr={"to": onnx_pb.TensorProto.FLOAT})
-    nodes.append(cast1)
     if axis == 1:
         mask_matrix = ctx.make_node(op_type="Transpose", inputs=cast1.output)
-        nodes.append(mask_matrix)
     else:
         mask_matrix = squeeze
     cast2 = ctx.make_node(op_type="Cast", inputs=mask_matrix.output,
                           attr={"to": ctx.get_dtype(node.input[0])})
-    nodes.append(cast2)
-    res = ctx.make_node(op_type="Mul", inputs=[cast2.output[0], node.input[0]],
-                        name=node.name, outputs=node.output)
-    nodes.append(res)
-    return nodes
+    shapes = node.output_shapes
+    dtypes = node.output_dtypes
+    ctx.remove_node(node.name)
+    ctx.make_node(op_type="Mul", inputs=[cast2.output[0], node.input[0]],
+                  name=node.name, outputs=node.output, shapes=shapes,
+                  dtypes=dtypes)

tf2onnx/function/range.py

@@ -11,7 +11,7 @@
 # pylint: disable=unused-argument,missing-docstring
 
 
-def make_range_const(ctx, start, limit, delta, output, scope_name, dtype):
+def make_range_const(ctx, start, limit, delta, output, scope_name, shape, dtype):
     """make Range subgraph if all inputs are const."""
     # T range = Range(T start, T limit, T delta)
     # V v_final_and_scan_outputs = Loop(int64 M, B cond, V v_initial)
@@ -21,59 +21,47 @@ def make_range_const(ctx, start, limit, delta, output, scope_name, dtype):
     delta = ctx.get_node_by_output(delta).get_tensor_value(as_list=False)
     val = np.arange(start, limit, delta, dtype=start.dtype)
     const_range = ctx.make_const(base_name, val)
-    return [ctx.make_node("Identity", [const_range.output[0]], dtypes=[dtype], outputs=[output]),
-            const_range]
+    ctx.make_node("Identity", [const_range.output[0]], shapes=[shape], dtypes=[dtype], outputs=[output])
 
 
-def make_range_non_const(ctx, start, limit, delta, output, scope_name, dtype):
+def make_range_non_const(ctx, start, limit, delta, output, scope_name, shape, dtype):
     """make Range subgraph."""
     # T range = Range(T start, T limit, T delta)
     # V v_final_and_scan_outputs = Loop(int64 M, B cond, V v_initial)
     base_name = utils.make_name(scope_name)
 
-    nodes = []
-
     # trip_count
     diff_node = ctx.make_node("Sub",
                               [limit, start],
                               op_name_scope=base_name,
                               name=utils.make_name("diff"))
     diff_output = diff_node.output[0]
-    nodes.append(diff_node)
 
     delta_cast = delta
     if dtype in [TensorProto.INT32, TensorProto.INT64]:
         cast_node = ctx.make_node("Cast", [diff_output], op_name_scope=base_name,
                                   name="cast_diff", attr={"to": TensorProto.FLOAT})
-        nodes.append(cast_node)
         diff_output = cast_node.output[0]
 
         cast_node = ctx.make_node("Cast", [delta], op_name_scope=base_name, name="cast_delta",
                                   attr={"to": TensorProto.FLOAT})
-        nodes.append(cast_node)
         delta_cast = cast_node.output[0]
-
     div_node = ctx.make_node("Div", [diff_output, delta_cast], op_name_scope=base_name, name="div")
-    nodes.append(div_node)
-
     ceil_node = ctx.make_node("Ceil", [div_node.output[0]], op_name_scope=base_name, name="ceil")
-    nodes.append(ceil_node)
-
     trip_count_node = ctx.make_node("Cast", [ceil_node.output[0]], op_name_scope=base_name, name="trip_cnt",
                                     attr={"to": TensorProto.INT64})
-    nodes.append(trip_count_node)
 
     # cond
     # Use initializer here since Constant OP before opset 9 does not support bool type
     cond_name = "{}_cond".format(base_name)
-    nodes.append(ctx.make_const(cond_name, np.ones((), dtype=bool)))
+    ctx.make_const(cond_name, np.ones((), dtype=bool))
 
     # body
     g = ctx.create_new_graph_with_same_config()
-    body_nodes = [g.make_node("Identity", ["cond"], outputs=["cond_out"]),
-                  g.make_node("Add", ["prev", delta], outputs=["current"], name=utils.make_name("add")),
-                  g.make_node("Identity", ["prev"], outputs=["range"])]
-    g.set_nodes(body_nodes)
+    g.make_node("Identity", ["cond"], outputs=["cond_out"])
+    g.make_node("Add", ["prev", delta], outputs=["current"], name=utils.make_name("add"))
+    g.make_node("Identity", ["prev"], outputs=["range"])
+
     g.add_graph_input("i", TensorProto.INT64, [])
     g.add_graph_input("cond", TensorProto.BOOL, [])
     g.add_graph_input("prev", dtype, [])
@@ -86,25 +74,25 @@ def make_range_non_const(ctx, start, limit, delta, output, scope_name, dtype):
     loop_inputs = [trip_count_node.output[0], cond_name, start]
     loop_node = ctx.make_node("Loop", loop_inputs, output_count=2, op_name_scope=base_name, name="loop")
     loop_node.set_body_graph_as_attr("body", g)
-    nodes.append(loop_node)
-
-    identity_node = ctx.make_node("Identity", [loop_node.output[1]], name=base_name, dtypes=[dtype], outputs=[output])
-    nodes.append(identity_node)
 
-    return nodes
+    ctx.make_node("Identity", [loop_node.output[1]], name=base_name, shapes=[shape],
+                  dtypes=[dtype], outputs=[output])
 
 
-def make_range(ctx, start, limit, delta, output, scope_name, dtype):
-    if all(ctx.get_node_by_output(n).is_const() for n in [start, limit, delta]):
-        return make_range_const(ctx, start, limit, delta, output, scope_name, dtype)
-    return make_range_non_const(ctx, start, limit, delta, output, scope_name, dtype)
+def make_range(ctx, start, limit, delta, output, scope_name, shape, dtype):
+    if all(ctx.get_node_by_output(n).is_const() for n in [start, limit, delta]) is True:
+        make_range_const(ctx, start, limit, delta, output, scope_name, shape, dtype)
+    else:
+        make_range_non_const(ctx, start, limit, delta, output, scope_name, shape, dtype)
 
 
 def range_op7(ctx, node, name, args):
     """Range."""
     # T range = Range(T start, T limit, T delta)
     # V v_final_and_scan_outputs = Loop(int64 M, B cond, V v_initial)
     dtype = node.get_attr_int("Tidx")
+    shape = node.output_shapes[0]
     utils.make_sure(dtype is not None, "Tidx of %s is None", node.name)
-    return make_range(ctx, node.input[0], node.input[1], node.input[2],
-                      node.output[0], name, dtype)
+    ctx.remove_node(node.name)
+    make_range(ctx, node.input[0], node.input[1], node.input[2],
+               node.output[0], name, shape, dtype)