Merge pull request #380 from onnx/gs/bert-regression

fix regression for bert, we broke if multiple model outputs came from the same node

Author: pengwa

tf2onnx/graph.py

@@ -349,8 +349,20 @@ 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

tf2onnx/tfonnx.py

@@ -250,11 +250,17 @@ def arg_minmax_op(ctx, node, name, args):
 
 def reduce_op(ctx, node, name, args):
     axes_node = node.inputs[1]
-    input_rank = len(ctx.get_shape(node.input[0]))
     axes = axes_node.get_tensor_value()
     if np.isscalar(axes):
         axes = [axes]
-    axes = [val + input_rank if val < 0 else val for val in axes]
+    input_shape = ctx.get_shape(node.input[0])
+    if input_shape is None:
+        if any([val < 0 for val in axes]):
+            raise ValueError("reduce_op: cannot have negative axis because we don't know input rank")
+    else:
+        input_rank = len(ctx.get_shape(node.input[0]))
+        axes = [val + input_rank if val < 0 else val for val in axes]
+
     node.set_attr("axes", axes)
     ctx.remove_input(node, node.input[1])
     keep_dims = node.get_attr("keep_dims")