Fix common bug where node inputs are assumed to come from output[0] (#1307)

Signed-off-by: Tom Wildenhain <[email protected]>

Author: TomWildenhain-Microsoft

tf2onnx/custom_opsets/ms.py

@@ -63,10 +63,10 @@ def version_1(cls, ctx, node, **kwargs):
         input_shape = ctx.make_node("Shape", [node.input[2]])
         hw_indices = ctx.make_const(utils.make_name("hw_indices"), np.array([1, 2]).astype(np.int64))
         input_shape_hw = ctx.make_node("Gather", [input_shape.output[0], hw_indices.output[0]])
-        output_shape = node.inputs[0]
-        if ctx.get_dtype(output_shape.output[0]) != onnx_pb.TensorProto.INT64:
-            output_shape = ctx.make_node("Cast", [output_shape.output[0]], attr={"to": onnx_pb.TensorProto.INT64})
-        output_shape_hw = ctx.make_node("Gather", [output_shape.output[0], hw_indices.output[0]])
+        output_shape = node.input[0]
+        if ctx.get_dtype(output_shape) != onnx_pb.TensorProto.INT64:
+            output_shape = ctx.make_node("Cast", [output_shape], attr={"to": onnx_pb.TensorProto.INT64}).output[0]
+        output_shape_hw = ctx.make_node("Gather", [output_shape, hw_indices.output[0]])
         kernel_shape_hw = list(ctx.get_shape(node.input[1]))[0:2]
         kernel_shape = ctx.make_const(utils.make_name("const_convtrans"), np.array(kernel_shape_hw).astype(np.int64))
         strides = conv_dims_attr(node, "strides")

tf2onnx/onnx_opset/math.py

@@ -154,25 +154,25 @@ def version_8(cls, ctx, node, **kwargs):
         # fetch those upfront since they are not accessible once we remove 'node'
         shapes = node.output_shapes
         dtypes = node.output_dtypes
-        input_dtype = node.inputs[0].output_dtypes[0]
+        input_dtype = ctx.get_dtype(node.input[0])
         name = node.name
-        min_node = node.inputs[1]
-        if min_node.output_dtypes[0] not in supported:
+        min_node = node.input[1]
+        if ctx.get_dtype(min_node) not in supported:
             # cast min if needed
-            min_node = ctx.insert_new_node_on_input(node, "Cast", min_node.output[0], to=onnx_pb.TensorProto.FLOAT)
-        max_node = node.inputs[2]
-        if max_node.output_dtypes[0] not in supported:
+            min_node = ctx.insert_new_node_on_input(node, "Cast", min_node, to=onnx_pb.TensorProto.FLOAT).output[0]
+        max_node = node.input[2]
+        if ctx.get_dtype(max_node) not in supported:
             # cast max if needed
-            max_node = ctx.insert_new_node_on_input(node, "Cast", max_node.output[0], to=onnx_pb.TensorProto.FLOAT)
+            max_node = ctx.insert_new_node_on_input(node, "Cast", max_node, to=onnx_pb.TensorProto.FLOAT).output[0]
         ctx.remove_node(name)
-        new_node = ctx.make_node("Max", [node.input[0], min_node.output[0]], outputs=[node.output[0]],
+        new_node = ctx.make_node("Max", [node.input[0], min_node], outputs=[node.output[0]],
                                  shapes=shapes, dtypes=dtypes)
         if input_dtype not in supported:
             # cast the data tensor if needed
             ctx.insert_new_node_on_input(new_node, "Cast", new_node.input[0], to=onnx_pb.TensorProto.FLOAT)
 
         new_node = ctx.insert_new_node_on_output("Min", new_node.output[0], name=utils.make_name(name))
-        new_node.input.append(max_node.output[0])
+        new_node.input.append(max_node)
         # copy shape and type
         ctx.set_dtype(new_node.output[0], dtypes[0])
         ctx.set_shape(new_node.output[0], shapes[0])

tf2onnx/onnx_opset/nn.py

@@ -315,10 +315,10 @@ def build_dynamic_target_size(ctx, transposed_intput, target_hw):
         A tensor of rank 2 containing [n c nh nw]
     """
     # We get the first half [n c] of the target shape
-    shape_of_transposed_input = ctx.make_node("Shape", [transposed_intput.output[0]])
+    shape_of_transposed_input = ctx.make_node("Shape", [transposed_intput])
     first_half_of_shape = GraphBuilder(ctx).make_slice(
         {"data": shape_of_transposed_input.output[0], "ends": [2], "starts": [0]})
-    target_size_int64 = ctx.make_node("Cast", [target_hw.output[0]], attr={'to': TensorProto.INT64})
+    target_size_int64 = ctx.make_node("Cast", [target_hw], attr={'to': TensorProto.INT64})
     # We build a tensor containing [n c nh nw]
     final_target_size = ctx.make_node("Concat", [first_half_of_shape, target_size_int64.output[0]], {'axis': 0})
     return final_target_size
@@ -916,10 +916,10 @@ def any_version_after11(cls, opset, ctx, node, **kwargs):
         mode = "nearest" if node.get_attr("method") is not None and node.get_attr(
             "method").s == b"nearest" else "linear"
         extrapolation_value = float(node.get_attr("extrapolation_value", "0").f)
-        input_x = node.inputs[0]
-        boxes = node.inputs[1]
-        box_ind = node.inputs[2]
-        crop_size = node.inputs[3]
+        input_x = node.input[0]
+        boxes = node.input[1]
+        box_ind = node.input[2]
+        crop_size = node.input[3]
         trip_name = utils.make_name(node.name + "_i")
         cond_name = utils.make_name(node.name + "_cond")
         cond_out_name = utils.make_name(node.name + "cond_out")
@@ -932,9 +932,9 @@ def any_version_after11(cls, opset, ctx, node, **kwargs):
         const_one = g.make_const(utils.make_name(node.name + "_const_one"), np.array([1], dtype=np.int32))
         const_one_long = g.make_const(utils.make_name(node.name + "_const_one_long"), np.array([1], dtype=np.int64))
         index_end = g.make_node("Add", [trip_name, const_one_long.output[0]])
-        box_index_from = g.make_node("Slice", [box_ind.output[0], trip_name, index_end.output[0]], name="Slice_a")
+        box_index_from = g.make_node("Slice", [box_ind, trip_name, index_end.output[0]], name="Slice_a")
         box_index_to = g.make_node("Add", [box_index_from.output[0], const_one.output[0]])
-        target_x = g.make_node("Slice", [input_x.output[0], box_index_from.output[0], box_index_to.output[0],
+        target_x = g.make_node("Slice", [input_x, box_index_from.output[0], box_index_to.output[0],
                                          const_zero.output[0]], name="Slice_b")
         transposed_x = g.make_node("Transpose", [target_x.output[0]], attr={'perm': constants.NHWC_TO_NCHW})
         const_zero_zero = g.make_const(utils.make_name(node.name + "_const_zero_zero"),
@@ -943,15 +943,15 @@ def any_version_after11(cls, opset, ctx, node, **kwargs):
                                      np.array([1, 1], dtype=np.float32))
         const_four = g.make_const(utils.make_name(node.name + "_const_four"), np.array([4], dtype=np.int64))
         const_empty_float = g.make_const(utils.make_name("const_empty_float"), np.array([], dtype=np.float32))
-        box = g.make_node("Slice", [boxes.output[0], trip_name, index_end.output[0], const_zero_long.output[0]],
+        box = g.make_node("Slice", [boxes, trip_name, index_end.output[0], const_zero_long.output[0]],
                           name="Slice_c")
         roi_raw = g.make_node("Reshape", [box.output[0], const_four.output[0]])
         roi_raw_first_half = GraphBuilder(g).make_slice({"data": roi_raw.output[0], "ends": [2], "starts": [0]})
         roi_raw_second_half = GraphBuilder(g).make_slice({"data": roi_raw.output[0], "ends": [4], "starts": [2]})
         roi_concat_1 = g.make_node("Concat", [const_zero_zero.output[0], roi_raw_first_half], attr={'axis': 0})
         roi_concat_2 = g.make_node("Concat", [const_one_one.output[0], roi_raw_second_half], attr={'axis': 0})
         final_roi = g.make_node("Concat", [roi_concat_1.output[0], roi_concat_2.output[0]], attr={'axis': 0})
-        final_crop_size = build_dynamic_target_size(g, transposed_x, crop_size)
+        final_crop_size = build_dynamic_target_size(g, transposed_x.output[0], crop_size)
         resized_x = g.make_node("Resize", [transposed_x.output[0], final_roi.output[0], const_empty_float.output[0],
                                            final_crop_size.output[0]],
                                 attr={"mode": mode, "extrapolation_value": extrapolation_value,
@@ -961,7 +961,7 @@ def any_version_after11(cls, opset, ctx, node, **kwargs):
         g.make_node("Identity", [cond_name], outputs=[cond_out_name])
         g.add_graph_output(cond_out_name, TensorProto.BOOL, [])
         g.add_graph_output(squeeze_x.output[0], ctx.get_dtype(node.input[0]), [-1, -1, -1])
-        trip_node = ctx.make_node("Size", [box_ind.output[0]])
+        trip_node = ctx.make_node("Size", [box_ind])
         cond_const = ctx.make_const(utils.make_name("cond"), np.ones((), dtype=np.bool))
         ctx.remove_node(node.name)
         branches = {"body": g}
@@ -1070,7 +1070,7 @@ def _convert_since_9(cls, ctx, node, op_type, use_target_size=False):
         # because onnxruntime only supports to scale the last two dims so transpose is inserted
         input_nchw = ctx.make_node("Transpose", [node.input[0]], {"perm": constants.NHWC_TO_NCHW})
         if use_target_size:
-            final_target_size = build_dynamic_target_size(ctx, input_nchw, node.inputs[1])
+            final_target_size = build_dynamic_target_size(ctx, input_nchw.output[0], node.input[1])
             roi = ctx.make_const(utils.make_name("roi"), np.array([]).astype(np.float32))
             const_empty_float = ctx.make_const(utils.make_name("const_empty_float"), np.array([], dtype=np.float32))
             resize_inputs = [