Merge pull request #300 from pengwa/fix_bing

Fix bugs found in real models.

Author: guschmue

tests/test_backend.py

@@ -1540,7 +1540,7 @@ def test_sparse_softmax_cross_entropy_with_logits_large_class(self):
 
         self._run_test_case([_OUTPUT], {_INPUT: label_val, _INPUT1: logits_val})
 
-    @unittest.skip("TODO: add a common utility for onnxruntime version check in another PR")
+    @unittest.skipIf(BACKEND in ["onnxruntime"], "onnxruntime Slice did not supported BOOL.")
     def test_matrix_band_part(self):
         input_val = np.random.randint(0, 666, (10, 15)).astype(np.int32)
         input_x = tf.placeholder(dtype=tf.int32, shape=[None, None], name=_TFINPUT)
@@ -1550,6 +1550,16 @@ def test_matrix_band_part(self):
         _ = tf.identity(res1, name=_TFOUTPUT1)
         self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: input_val})
 
+    @unittest.skipIf(BACKEND in ["onnxruntime"], "onnxruntime Slice did not supported BOOL.")
+    def test_matrix_band_part_2(self):
+        input_val = np.random.randint(0, 666, (1, 1)).astype(np.int32)
+        input_x = tf.placeholder(dtype=tf.int32, shape=[None, None], name=_TFINPUT)
+        res = tf.matrix_band_part(input_x, -1, 0)
+        res1 = tf.matrix_band_part(input_x, 0, -1)
+        _ = tf.identity(res, name=_TFOUTPUT)
+        _ = tf.identity(res1, name=_TFOUTPUT1)
+        self._run_test_case([_OUTPUT, _OUTPUT1], {_INPUT: input_val})
+
     def test_floordiv(self):
         input_val_1 = np.random.random_sample(100).astype(np.int32)
         input_val_2 = (np.random.random_sample(100) + 1).astype(np.int32)

tf2onnx/function/matrixbandpart.py

@@ -41,13 +41,18 @@ def matrixbandpart_op(ctx, node, name, args):
     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))
-    slice_node = g.make_node(op_type="Slice", inputs=["line"],
+    ctx.set_dtype(const_zero_bool.output[0], onnx_pb.TensorProto.BOOL)
+
+    # shift right the line and add zero at the left.
+    new_line = g.make_node(op_type="Concat", inputs=[const_zero_bool.output[0], "line"], attr={"axis": counter_axis},
+                           dtypes=[onnx_pb.TensorProto.BOOL])
+    slice_node = g.make_node(op_type="Slice", inputs=[new_line.output[0]],
                              attr={"axes": [counter_axis], "starts": [0], "ends": [-1]})
-    new_line = g.make_node(op_type="Concat", inputs=[const_zero_bool.output[0], slice_node.output[0]],
-                           outputs=["line_out"], attr={"axis": counter_axis})
+
     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", ["line"], outputs=["res"]),
+                  g.make_node("Identity", [slice_node.output[0]], outputs=["line_out"])]
     g.set_nodes(body_nodes)
     g.add_graph_input("trip", onnx_pb.TensorProto.INT64, [])
     g.add_graph_input("cond", onnx_pb.TensorProto.BOOL, [])

tf2onnx/function/select.py

@@ -204,7 +204,9 @@ def create_loop_body_graph(parent_g, gather_input_ids, output_data_type, output_
 
     g.add_graph_output(cond_output_id, TensorProto.BOOL, ())
     g.add_graph_output(fake_var_output_id, TensorProto.FLOAT, ())
-    g.add_graph_output(loop_output_id, output_data_type, output_shape[1:])
+
+    # use None for all dims, just keep original rank. Because it is observed, dims might be changed in loop.
+    g.add_graph_output(loop_output_id, output_data_type, utils.create_vague_shape_like(output_shape[1:]))
 
     return g
 
@@ -234,5 +236,5 @@ def create_body_graph_for_if_branch(parent_g, data_type, output_shape, chosen_cu
     )
     nodes.append(identity_node)
     g.set_nodes(nodes)
-    g.add_graph_output("y", data_type, output_shape)
+    g.add_graph_output("y", data_type, utils.create_vague_shape_like(output_shape))
     return g

tf2onnx/graph.py

@@ -940,9 +940,9 @@ def _extract_sub_graph_nodes(self, dest_node, input_checker=None):
                 # we don't care about nested graph here, just handle current graph cropping.
                 node = self.get_node_by_output(input_id, search_in_parent_graphs=False)
                 if not node:
-                    # some node (for example Scan) has optional inputs, which
-                    # might has empty input.
-                    # subgraph might has input defined in outer graph
+                    # some nodes (for example Scan) have optional inputs, which
+                    # might have empty input.
+                    # subgraph might have input defined in outer graph
                     continue
                 if node not in res_set:
                     if input_checker and input_checker(node) is False:

tf2onnx/rewriter/loop_rewriter.py

@@ -56,6 +56,9 @@ def rewrite(self, context):
 
             body_nodes = set(cell_g_info.nodes + cond_g_info.nodes)
             body_outputs = cond_g_info.outputs + cell_g_info.outputs
+            for out_tensor_value_info in body_outputs:
+                out_tensor_value_info.shape = utils.create_vague_shape_like(out_tensor_value_info.shape)
+
             loop_body_g = LoopRewriterBase.construct_graph_from_nodes(self.g, body_nodes, body_outputs)
 
             # create loop body graph inputs
@@ -74,7 +77,7 @@ def rewrite(self, context):
                     dtype = tensor_value_info.dtype
                     shape = tensor_value_info.shape
 
-                loop_body_g.add_graph_input(input_name, dtype, shape)
+                loop_body_g.add_graph_input(input_name, dtype, utils.create_vague_shape_like(shape))
 
             body_nodes_to_append = []
             for input_ta in loop_props.tensor_array_inputs:

tf2onnx/tfonnx.py

@@ -1300,17 +1300,10 @@ def minmax_op(ctx, node, name, args):
 
 def pack_op(ctx, node, name, args):
     # hack to make up for the missing onnx pack op
-
-    pack_shape = ctx.get_shape(node.output[0])
-    if not pack_shape:
-        # sometimes Pack output shape is None (for example Pack is following control flow Exit op)
-        input_cnt = len(node.inputs)
-        input_shape = ctx.get_shape(node.input[0])
-        if input_shape:
-            pack_shape = [input_cnt] + input_shape
-            ctx.set_shape(node.output[0], pack_shape)
-
     axis = node.get_attr("axis").i
+    if axis < 0:
+        axis += len(ctx.get_shape(node.input[0])) + 1
+
     nodes = []
     inputs = []
     dtype = None
@@ -1747,6 +1740,14 @@ def zeroslike_op(ctx, node, name, args):
     return mul_op
 
 
+def softmax_op(ctx, node, name, args):
+    # T output = Softmax(T logits). The axis softmax would be performed on is always on -1.
+    # T output = Softmax(T input, @int axis). Default axis is 1.
+    logits_rank = len(ctx.get_shape(node.input[0]))
+    node.set_attr("axis", logits_rank - 1)
+    return node
+
+
 # map tensorflow ops to onnx ops. The format below is
 # "TFOP": func_to_map, ["OnnxOp", ...]
 #
@@ -1829,7 +1830,7 @@ def zeroslike_op(ctx, node, name, args):
     "Sqrt": (direct_op, []),
     "Square": (square_op, []),
     "SquaredDifference": (squareddifference_op, []),
-    "Softmax": (direct_op, ["Softmax"]),
+    "Softmax": (softmax_op, ["Softmax"]),
     "StopGradient": (identity_op, ["Identity"]),
     "StridedSlice": (stridedslice_op, []),
     "Sub": (broadcast_op, []),
@@ -2261,6 +2262,10 @@ def rewrite_incomplete_type_support(g, ops, impacted_ops):
 
                 input_name = op.input[i]
                 dtype = g.get_dtype(input_name)
+                if dtype is None:
+                    log.warning("Adding Cast for op %s (type is %s)' input: %s, dtype should not be None",
+                                op.name, op.type, input_name)
+
                 if dtype != onnx_pb.TensorProto.FLOAT:
                     output_dtype = dtype
                     if input_node and input_node.type == "Cast" \

tf2onnx/utils.py

@@ -335,3 +335,8 @@ def construct_graph_from_nodes(parent_g, nodes, outputs, shapes, dtypes):
 
 def tf_name_scope(name):
     return '/'.join(name.split('/')[:-1])
+
+
+def create_vague_shape_like(shape):
+    make_sure(len(shape) >= 0, "rank should be >= 0")
+    return [-1 for i in enumerate(shape)]