Merge pull request #397 from pengwa/fix_issue_

fix an issue when subgraph output node has more than one output

Author: guschmue

tests/test_backend.py

@@ -883,6 +883,12 @@ def test_split(self):
         _ = tf.identity(x_, name=_TFOUTPUT)
         self._run_test_case([_OUTPUT], {_INPUT: x_val})
 
+    def test_split_with_more_outputs(self):
+        x_val = np.linspace(1.0, 5 * 30.0, 5 * 30).astype(np.float32).reshape(5, 30)
+        x0 = tf.placeholder(tf.float32, x_val.shape, name=_TFINPUT)
+        _, _, _ = tf.split(x0, [4, 15, 11], 1, name="split_test")
+        self._run_test_case(["split_test:0", "split_test:1", "split_test:2"], {_INPUT: x_val})
+
     def test_reducesum(self):
         # not supported by onnx-caffe2
         x_val = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32).reshape((2, 2))

tf2onnx/graph.py

@@ -357,27 +357,15 @@ def __init__(self, nodes, output_shapes=None, dtypes=None, target=None, opset=No
         self.reset_nodes(ops)
 
         # add identity node after each output, in case it is renamed during conversion.
-        nodes_seen = set()
-        multi_output_nodes = set()
         for o in self.outputs:
             n = self.get_node_by_output_in_current_graph(o)
-            if n in nodes_seen:
-                multi_output_nodes.add(n)
-            else:
-                nodes_seen.add(n)
-
-        for o in self.outputs:
-            n = self.get_node_by_output_in_current_graph(o)
-            # TODO: below doesn't work for nodes with multiple outputs. A work around, keep those intact.
-            if n in multi_output_nodes:
-                continue
             new_output_name = port_name(n.name + "_" + utils.make_name("raw_output_"))
             n_shapes = n.output_shapes
             n_dtypes = n.output_dtypes
             body_graphs = n.graph.contained_graphs.pop(n.name, None)
             self.remove_node(n.name)
 
-            new_outputs = [o if o != output else new_output_name for output in n.output]
+            new_outputs = [output if output != o else new_output_name for output in n.output]
             # domain should be passed to new node
             new_node = self.make_node(n.type, n.input, outputs=new_outputs, attr=n.attr, name=n.name,
                                       skip_conversion=n._skip_conversion, dtypes=n_dtypes, shapes=n_shapes,