Tom/tflite flexops (#1374)

* Fix parsing of tflite string tensors and other tflite fixes

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

* WIP

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

* Work around string decoding bug for flex ops

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

* Topsort tflite subgraphs

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

* Disable some tflite tests

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

* Pylint

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

Author: TomWildenhain-Microsoft

tests/backend_test_base.py

@@ -197,6 +197,10 @@ def convert_to_tflite(self, graph_def, feed_dict, outputs):
             sess_outputs = [sess.graph.get_tensor_by_name(n) for n in outputs]
             converter = tf_lite.TFLiteConverter.from_session(sess, sess_inputs, sess_outputs)
             #converter.optimizations = [tf.lite.Optimize.DEFAULT]
+            converter.target_spec.supported_ops = [
+                tf.lite.OpsSet.TFLITE_BUILTINS,    # enable TensorFlow Lite ops.
+                tf.lite.OpsSet.SELECT_TF_OPS,      # enable TensorFlow flex ops.
+            ]
 
             from tensorflow.lite.python.convert import ConverterError
             try:

tests/test_backend.py

@@ -1531,6 +1531,7 @@ def func(data, segments):
             self._run_test_case(func, [_OUTPUT], {_INPUT: data_val, _INPUT1: segs_val})
 
     @check_opset_min_version(11, "Pad")
+    @skip_tflite("unknown rank")
     def test_segment_mean_unknown_rank(self):
         segs_val = np.array([0, 0, 0, 1, 2, 2, 3, 3], dtype=np.int32)
         data_val = np.arange(8 * 2 * 3, dtype=np.float32).reshape([8, 2, 3])
@@ -1820,7 +1821,7 @@ def func():
             return tf.identity(x_, name=_TFOUTPUT)
         # since results are random, compare the shapes only
         g = self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
-        results = self.run_backend(g, [_OUTPUT], {})
+        results = self.run_backend(g, g.outputs, {})
         numbers = set(results[0].flatten())
         self.assertEqual(sorted(numbers), list(range(2, 10)))
 
@@ -1833,7 +1834,7 @@ def func():
             return tf.identity(x_, name=_TFOUTPUT)
         # since results are random, compare the shapes only
         g = self._run_test_case(func, [_OUTPUT], {}, check_value=False, check_shape=True)
-        results = self.run_backend(g, [_OUTPUT], {})
+        results = self.run_backend(g, g.outputs, {})
         self.assertTrue(2 <= results[0] < 10)
 
     def test_randomuniform_int_nonconst_max(self):
@@ -1845,7 +1846,11 @@ def func(m):
             x_ = tf.identity(x_, name="output2")
             return tf.identity(x_, name=_TFOUTPUT)
         g = self._run_test_case(func, [_OUTPUT], {_INPUT: m_val}, check_value=False, check_shape=True)
-        results = self.run_backend(g, [_OUTPUT], {_INPUT: m_val})
+        feed_dict = {_INPUT: m_val}
+        if "input" in g.input_names:
+            # TFLite inputs don't have port numbers
+            feed_dict = {k.split(":")[0]: v for k, v in feed_dict.items()}
+        results = self.run_backend(g, g.outputs, feed_dict)
         numbers = set(results[0].flatten())
         self.assertEqual(sorted(numbers), list(range(8)))
 
@@ -1859,7 +1864,11 @@ def func(n, m):
             x_ = tf.identity(x_, name="output2")
             return tf.identity(x_, name=_TFOUTPUT)
         g = self._run_test_case(func, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val}, check_value=False, check_shape=True)
-        results = self.run_backend(g, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val})
+        feed_dict = {_INPUT: n_val, _INPUT1: m_val}
+        if "input" in g.input_names:
+            # TFLite inputs don't have port numbers
+            feed_dict = {k.split(":")[0]: v for k, v in feed_dict.items()}
+        results = self.run_backend(g, g.outputs, feed_dict)
         numbers = set(results[0].flatten())
         self.assertEqual(sorted(numbers), list(range(2, 10)))
 
@@ -1875,7 +1884,11 @@ def func(n, m, s):
             return tf.identity(x_, name=_TFOUTPUT)
         g = self._run_test_case(func, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val, _INPUT2: s_val},
                                 check_value=False, check_shape=True)
-        results = self.run_backend(g, [_OUTPUT], {_INPUT: n_val, _INPUT1: m_val, _INPUT2: s_val})
+        feed_dict = {_INPUT: n_val, _INPUT1: m_val, _INPUT2: s_val}
+        if "input" in g.input_names:
+            # TFLite inputs don't have port numbers
+            feed_dict = {k.split(":")[0]: v for k, v in feed_dict.items()}
+        results = self.run_backend(g, g.outputs, feed_dict)
         numbers = set(results[0].flatten())
         self.assertEqual(sorted(numbers), list(range(2, 10)))
 
@@ -4097,6 +4110,7 @@ def func(splits, rt_dense_values, indices):
 
     @check_tf_min_version("1.14", "ragged needs tf 1.14")
     @check_opset_min_version(11, "CumSum")
+    @skip_tflite("unknown rank")
     def test_ragged_tensor_to_tensor(self):
         splits_val1 = np.array([0, 1, 1, 5], dtype=np.int32)
         splits_val2 = np.array([0, 3, 3, 5, 9, 10], dtype=np.int32)
@@ -4758,6 +4772,7 @@ def func(input_val):
             self.config.opset = current_opset
 
     @check_tf_min_version("1.14")
+    @skip_tflite("FlexRFFT2D")
     def test_rfft_ops(self):
 
         def dft_slow(x, M):

tests/test_tflite_postprocess.py

@@ -167,7 +167,9 @@ def make_postprocess_model(self, max_detections=10, detections_per_class=100, ma
         operators = builder.EndVector(1)
 
         # subgraph
+        graph_name = builder.CreateString("TFLite graph")
         SubGraph.SubGraphStart(builder)
+        SubGraph.SubGraphAddName(builder, graph_name)
         SubGraph.SubGraphAddTensors(builder, tensors)
         SubGraph.SubGraphAddInputs(builder, inputs)
         SubGraph.SubGraphAddOutputs(builder, outputs)

tests/test_tflite_utils.py

@@ -59,9 +59,9 @@ def func(a, b, c):
         self.assertEqual(1, len(tflite_graphs))
         onnx_nodes, op_cnt, attr_cnt, output_shapes, dtypes, inputs, outputs, _ = \
             parse_tflite_graph(tflite_graphs[0], opcodes_map, model, tensor_shapes_override=tensor_shapes)
-        self.assertEqual(2, op_cnt['MUL'])
-        self.assertEqual(1, op_cnt['ADD'])
-        self.assertEqual(1, op_cnt['FULLY_CONNECTED'])
+        self.assertEqual(2, op_cnt['TFL_MUL'])
+        self.assertEqual(1, op_cnt['TFL_ADD'])
+        self.assertEqual(1, op_cnt['TFL_FULLY_CONNECTED'])
 
         self.assertEqual(1, attr_cnt['WeightsFormat'])
         self.assertEqual(names, inputs)

tf2onnx/flexbuffers.py

@@ -32,8 +32,12 @@ def read_float(buffer, offset, bit_size):
     raise FlexbufferParseException("Invalid bit size for flexbuffer float: %d" % bit_size)
 
 
-def read_string(buffer, offset, size):
-    return buffer[offset:offset+size].decode('utf-8')
+def read_string(buffer, offset, size, decode_strings):
+    data = buffer[offset:offset+size]
+    if decode_strings:
+        # Flexbuffer requires all strings to be valid UTF-8 but FlexOps don't always respect this.
+        data = data.decode('utf-8')
+    return data
 
 
 def read_indirect(buffer, offset, bit_size):
@@ -44,16 +48,16 @@ def read_bytes(buffer, offset, size):
     return buffer[offset:offset+size]
 
 
-def read_array(buffer, offset, length, bit_size, packed_type):
+def read_array(buffer, offset, length, bit_size, packed_type, decode_strings):
     byte_size = 1 << bit_size
     arr = []
     for i in range(length):
         item_offset = offset + (i * byte_size)
-        arr.append(read_buffer(buffer, item_offset, bit_size, packed_type))
+        arr.append(read_buffer(buffer, item_offset, bit_size, packed_type, decode_strings))
     return arr
 
 
-def read_buffer(buffer, offset, parent_bit_size, packed_type):
+def read_buffer(buffer, offset, parent_bit_size, packed_type, decode_strings):
     """Recursively decode flatbuffer object into python representation"""
     bit_size = packed_type & 3
     value_type = packed_type >> 2
@@ -64,20 +68,22 @@ def read_buffer(buffer, offset, parent_bit_size, packed_type):
     if value_type in [0x1, 0x2, 0x3]:
         read_fn = {0x1: read_int, 0x2: read_uint, 0x3: read_float}[value_type]
         return read_fn(buffer, offset, parent_bit_size)
-    if value_type in [0x4, 0x5]:
+    if value_type == 0x4:
         str_offset = read_indirect(buffer, offset, parent_bit_size)
         size = 0
         while read_int(buffer, str_offset + size, 0) != 0:
             size += 1
-        return read_string(buffer, str_offset, size)
+        return read_string(buffer, str_offset, size, decode_strings)
     if value_type == 0x5:
         str_offset = read_indirect(buffer, offset, parent_bit_size)
-        size_byte_size = 1 << bit_size
+        size_bit_size = bit_size
+        size_byte_size = 1 << size_bit_size
         size = read_uint(buffer, str_offset - size_byte_size, bit_size)
         while read_int(buffer, str_offset + size, 0) != 0:
             size_byte_size <<= 1
-            size = read_uint(buffer, str_offset - size_byte_size, bit_size)
-        return read_string(buffer, str_offset, size)
+            size_bit_size += 1
+            size = read_uint(buffer, str_offset - size_byte_size, size_bit_size)
+        return read_string(buffer, str_offset, size, decode_strings)
     if value_type in [0x6, 0x7, 0x8]:
         read_fn = {0x6: read_int, 0x7: read_uint, 0x8: read_float}[value_type]
         data_offset = read_indirect(buffer, offset, parent_bit_size)
@@ -93,10 +99,10 @@ def read_buffer(buffer, offset, parent_bit_size, packed_type):
         obj = {}
         for i in range(length):
             key_offset = keys_vector_offset + i * key_byte_size
-            key = read_buffer(buffer, key_offset, key_bit_size, (0x4 << 2) | key_bit_size)
+            key = read_buffer(buffer, key_offset, key_bit_size, (0x4 << 2) | key_bit_size, decode_strings)
             value_offset = values_offset + i * byte_size
             value_packed_type = read_uint(buffer, packed_types_offset + i, 0)
-            value = read_buffer(buffer, value_offset, bit_size, value_packed_type)
+            value = read_buffer(buffer, value_offset, bit_size, value_packed_type, decode_strings)
             obj[key] = value
         return obj
     if value_type == 0xa:
@@ -107,21 +113,21 @@ def read_buffer(buffer, offset, parent_bit_size, packed_type):
         for i in range(length):
             item_offset = items_offset + (i * byte_size)
             packed_type = read_uint(buffer, packed_types_offset + i, 0)
-            arr.append(read_buffer(buffer, item_offset, bit_size, packed_type))
+            arr.append(read_buffer(buffer, item_offset, bit_size, packed_type, decode_strings))
         return arr
     if value_type in [0xb, 0xc, 0xd, 0xe, 0xf, 0x24]:
         length_offset = read_indirect(buffer, offset, parent_bit_size) - byte_size
         length = read_uint(buffer, length_offset, bit_size)
         item_value_type = value_type - 0xb + 0x1
         packed_type = item_value_type << 2
         items_offset = read_indirect(buffer, offset, parent_bit_size)
-        return read_array(buffer, items_offset, length, bit_size, packed_type)
+        return read_array(buffer, items_offset, length, bit_size, packed_type, decode_strings)
     if 0x10 <= value_type <= 0x18:
         length = (value_type - 0x10) // 3 + 2
         value_type = ((value_type - 0x10) % 3) + 1
         packed_type = value_type << 2
         items_offset = read_indirect(buffer, offset, parent_bit_size)
-        return read_array(buffer, items_offset, length, bit_size, packed_type)
+        return read_array(buffer, items_offset, length, bit_size, packed_type, decode_strings)
     if value_type == 0x19:
         data_offset = read_indirect(buffer, offset, parent_bit_size)
         size_offset = data_offset - byte_size
@@ -132,9 +138,9 @@ def read_buffer(buffer, offset, parent_bit_size, packed_type):
     raise FlexbufferParseException("Invalid flexbuffer value type %r" % value_type)
 
 
-def read_flexbuffer(buffer):
+def read_flexbuffer(buffer, decode_strings=True):
     byte_size = read_uint(buffer, len(buffer) - 1, 0)
     bit_size = {1: 0, 2: 1, 4: 2, 8: 3, 16: 4}[byte_size]
     packed_type = read_uint(buffer, len(buffer) - 2, 0)
     offset = len(buffer) - 2 - byte_size
-    return read_buffer(buffer, offset, bit_size, packed_type)
+    return read_buffer(buffer, offset, bit_size, packed_type, decode_strings)

tf2onnx/onnx_opset/math.py

@@ -584,6 +584,7 @@ def version_10(cls, ctx, node, **kwargs):
 
         shapes = node.output_shapes
         dtypes = [onnx_pb.TensorProto.BOOL] * len(node.output_dtypes)
+        outputs = node.output
 
         ctx.remove_node(node.name)
 
@@ -593,7 +594,7 @@ def version_10(cls, ctx, node, **kwargs):
                                  shapes=shapes, dtypes=dtypes)
         or_node = ctx.make_node("Or", inputs=[inf_node.output[0], nan_node.output[0]], name=utils.make_name(node.name),
                                 shapes=shapes, dtypes=dtypes)
-        _ = ctx.make_node("Not", inputs=or_node.output, name=node.name,
+        _ = ctx.make_node("Not", inputs=or_node.output, name=node.name, outputs=outputs,
                           shapes=shapes, dtypes=dtypes)
 
 

tf2onnx/tf_utils.py

@@ -318,26 +318,75 @@ def replace_placeholders_with_tables(graph_def, placeholder_to_table_info):
             n.attr['key_dtype'].type = key_dtype
             n.attr['value_dtype'].type = val_dtype
 
-def tflist_to_onnx(g, shape_override, const_node_values=None, ignore_default=None, use_default=None):
-    """
-    Convert the tf-node list into an onnx graph with minimal rewrites so
-    we can use the onnx graph as intermediate graph.
-    """
+def read_tf_node_def_attrs(node_def, input_dtypes, input_shapes):
+    """Given a tf node def, returns a dict of attribute names to values"""
+    from tf2onnx.tf_loader import tf_session, tf_placeholder  # pylint: disable=import-outside-toplevel
+    del node_def.input[:]
+    node_def.name = "node"
+
+    # read_tf_node_attrs uses some tf methods that require the node to be loaded into a valid TF graph
+    g = tf.Graph()
+    with g.as_default():
+        for i, (dtype, shape) in enumerate(zip(input_dtypes, input_shapes)):
+            inp = "input" + str(i)
+            tf_placeholder(dtype, name=inp, shape=shape)
+            node_def.input.append(inp)
+        mini_graph_def = g.as_graph_def()
+        mini_graph_def.node.append(node_def)
+    g2 = tf.Graph()
+    with g2.as_default():
+        with tf_session() as sess:
+            tf.import_graph_def(mini_graph_def, name='')
+            node = sess.graph.get_operation_by_name("node")
+            return read_tf_node_attrs(node)
+
+
+def read_tf_node_attrs(node):
+    """Given a tf Node, returns a dict of attribute names to values"""
+    attr = {}
+    attr_cnt = collections.Counter()
 
     # ignore the following attributes
-    ignored_attr = {"unknown_rank", "_class", "Tshape", "use_cudnn_on_gpu", "Index", "Tpaddings",
+    ignored_attr = {"T", "unknown_rank", "_class", "Tshape", "use_cudnn_on_gpu", "Index", "Tpaddings",
                     "TI", "Tparams", "Tindices", "Tlen", "Tdim", "Tin", "dynamic_size", "Tmultiples",
                     "Tblock_shape", "Tcrops", "index_type", "Taxis", "U", "maxval",
                     "Tout", "Tlabels", "Tindex", "element_shape", "Targmax", "Tperm", "Tcond",
                     "T_threshold", "element_dtype", "shape_type", "_lower_using_switch_merge",
                     "parallel_iterations", "_num_original_outputs", "output_types", "output_shapes",
                     "key_dtype", "value_dtype", "Tin", "Tout", "capacity", "component_types", "shapes",
                     "Toutput_types", "dense_shapes", "Tdense", "Tsegmentids", "Tshift", "Tnumsegments", "SrcT",
+                    "body", "cond", "then_branch", "else_branch", "f",
                     "Tcomplex", "Treal",  # For RFFT, Tcomplex is ignored because
                                           # onnx.helper.make_node fails,
                                           # TODO: it should be added back.
                     }
 
+    for a in node.node_def.attr:
+        attr_cnt[a] += 1
+        value = get_tf_node_attr(node, a)
+        if a in ignored_attr or isinstance(value, tensor_pb2.TensorProto):
+            pass
+        elif a == "shape":
+            shape = get_tf_shape_attr(node)
+            if shape is not None:
+                attr[a] = shape
+        elif a == "DstT":
+            attr["to"] = map_tf_dtype(value)
+        elif isinstance(value, tf.DType):
+            attr[a] = map_tf_dtype(value)
+        elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], tf.DType):
+            attr[a] = [map_tf_dtype(v) for v in value]
+        else:
+            attr[a] = get_tf_node_attr(node, a)
+
+    return attr, attr_cnt
+
+def tflist_to_onnx(g, shape_override, const_node_values=None, ignore_default=None, use_default=None):
+    """
+    Convert the tf-node list into an onnx graph with minimal rewrites so
+    we can use the onnx graph as intermediate graph.
+    """
+
     node_list = g.get_operations()
     functions = {}
 
@@ -360,41 +409,27 @@ def tflist_to_onnx(g, shape_override, const_node_values=None, ignore_default=Non
             dtypes[out.name] = map_tf_dtype(out.dtype)
             output_shapes[out.name] = shape
 
-    # minimal conversion of attributes
     for node in ops:
-        attr = {}
+        attr, new_attr_cnt = read_tf_node_attrs(node)
+        attr_cnt += new_attr_cnt
         takeit = True
         op_cnt[node.type] += 1
         for a in node.node_def.attr:
             attr_cnt[a] += 1
             value = get_tf_node_attr(node, a)
-            if a in ignored_attr:
-                pass
-            elif a == "T":
+            if a == "T":
                 if value and not isinstance(value, list):
                     dtypes[node.name] = map_tf_dtype(value)
-            elif a == "shape":
-                shape = get_tf_shape_attr(node)
-                if shape is not None:
-                    attr[a] = shape
             elif a in {"body", "cond", "then_branch", "else_branch", "f"}:
                 input_shapes = [inp.get_shape() for inp in node.inputs]
                 nattr = get_tf_node_attr(node, a)
                 attr[a] = nattr.name
                 functions[nattr.name] = input_shapes
-            elif a == "DstT":
-                attr["to"] = map_tf_dtype(value)
             elif isinstance(value, tensor_pb2.TensorProto):
                 if const_node_values and node.name in const_node_values:
                     value.tensor_content = const_node_values[node.name]
                 onnx_tensor = tf_to_onnx_tensor(value, name=port_name(node.name))
                 attr[a] = onnx_tensor
-            elif isinstance(value, tf.DType):
-                attr[a] = map_tf_dtype(value)
-            elif isinstance(value, list) and len(value) > 0 and isinstance(value[0], tf.DType):
-                attr[a] = [map_tf_dtype(v) for v in value]
-            else:
-                attr[a] = get_tf_node_attr(node, a)
 
         node_type = node.type
         input_names = [i.name for i in node.inputs]

tf2onnx/tflite_handlers/tfl_math.py

@@ -66,6 +66,7 @@ def to_tf(cls, ctx, node, **kwargs):
         pass
 
 @tfl_op(["TFL_REDUCE_MAX"], tf_op="Max")
+@tfl_op(["TFL_REDUCE_MIN"], tf_op="Min")
 @tfl_op(["TFL_REDUCE_ANY"], tf_op="Any")
 @tfl_op(["TFL_REDUCE_PROD"], tf_op="Prod")
 class TflReduceOp: