[TFLite] Support TFLite FP32 Relay frontend. (#2365)

* Support TFLite FP32 Relay frontend.

* Fix lint issue

* Remove unnecessary variables and packages

* Add method doc string

* Fix capital letter of method doc string

* Add TFLite FP32 Relay frontend based on latest code

* Merge the latest code

* Solve cython issue with relay type inference

* Modify code based on suggestions

python/tvm/relay/frontend/__init__.py

@@ -10,3 +10,4 @@ from __future__ import absolute_import
 from .mxnet import from_mxnet
 from .keras import from_keras
 from .onnx import from_onnx
+from .tflite import from_tflite

tests/python/frontend/tflite/test_forward.py

+# pylint: disable=import-self, invalid-name, unused-argument
+"""
+TFLite testcases
+================
+This article is a test script to test TFLite operator with Relay.
+"""
+from __future__ import print_function
+import numpy as np
+import tvm
+from tvm import relay
+from tvm.contrib import util
+import tensorflow as tf
+from tensorflow.python.framework import constant_op
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.ops import array_ops
+from tensorflow.python.ops import variables
+from tensorflow.contrib.lite.python import interpreter as interpreter_wrapper
+
+import nnvm.testing.tf
+
+#######################################################################
+# Generic run functions for TVM & TFLite
+# --------------------------------------
+def convert_to_list(x):
+    if not isinstance(x, list):
+        x = [x]
+    return x
+
+def run_tvm_graph(tflite_model_buf, input_data, input_node, num_output=1, target='llvm',
+                  out_names=None):
+    """ Generic function to compile on relay and execute on tvm """
+    try:
+        import tflite.Model
+    except ImportError:
+        raise ImportError("The tflite package must be installed")
+
+    # get TFLite model from buffer
+    tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)
+
+    input_data = convert_to_list(input_data)
+    input_node = convert_to_list(input_node)
+
+    shape_dict = {}
+    dtype_dict = {}
+    for i, e in enumerate(input_node):
+        shape_dict[e] = input_data[i].shape
+        dtype_dict[e] = input_data[i].dtype.name
+
+    func, params = relay.frontend.from_tflite(tflite_model,
+                                              shape_dict=shape_dict,
+                                              dtype_dict=dtype_dict)
+    with relay.build_config(opt_level=3):
+        graph, lib, params = relay.build(func, target, params=params)
+
+    ctx = tvm.context(target, 0)
+    from tvm.contrib import graph_runtime
+    m = graph_runtime.create(graph, lib, ctx)
+    # set inputs
+    for i, e in enumerate(input_node):
+        m.set_input(e, tvm.nd.array(input_data[i].astype(input_data[i].dtype)))
+
+    m.set_input(**params)
+    # execute
+    m.run()
+    # get outputs
+    assert out_names is None or num_output == len(out_names), "out_names: {} num_output: {}".format(
+        out_names, num_output)
+    tvm_output_list = []
+    for i in range(0, num_output):
+        tvm_output = m.get_output(i)
+        tvm_output_list.append(tvm_output.asnumpy())
+    return tvm_output_list
+
+
+def run_tflite_graph(tflite_model_buf, input_data):
+    """ Generic function to execute TFLite """
+    input_data = convert_to_list(input_data)
+
+    interpreter = interpreter_wrapper.Interpreter(model_content=tflite_model_buf)
+    interpreter.allocate_tensors()
+
+    input_details = interpreter.get_input_details()
+    output_details = interpreter.get_output_details()
+
+    # set input
+    assert len(input_data) == len(input_details)
+    for i in range(len(input_details)):
+        interpreter.set_tensor(input_details[i]['index'], input_data[i])
+
+    # Run
+    interpreter.invoke()
+
+    # get output
+    tflite_output = list()
+    for i in range(len(output_details)):
+        tflite_output.append(interpreter.get_tensor(output_details[i]['index']))
+
+    return tflite_output
+
+
+def compare_tflite_with_tvm(tflite_in_data, tvm_in_data, in_name, input_tensors,
+                            output_tensors, output_need_transpose_nchw=False,
+                            init_global_variables=False):
+    """Generic function to generate and compare TFLite and TVM output"""
+    tflite_in_data = convert_to_list(tflite_in_data)
+    tvm_in_data = convert_to_list(tvm_in_data)
+    in_name = convert_to_list(in_name)
+    in_node = [0] * len(in_name)
+    for i in range(len(in_name)):
+        in_node[i] = in_name[i].split(':')[0] if ":" in in_name[i] else in_name[i]
+
+    with tf.Session() as sess:
+        if init_global_variables:
+            sess.run(variables.global_variables_initializer())
+        # convert to tflite model
+        converter = tf.contrib.lite.TFLiteConverter.from_session(
+            sess, input_tensors, output_tensors)
+        tflite_model_buffer = converter.convert()
+        tflite_output = run_tflite_graph(tflite_model_buffer, tflite_in_data)
+
+        for device in ["llvm"]:
+            ctx = tvm.context(device, 0)
+            if not ctx.exist:
+                print("Skip because %s is not enabled" % device)
+                continue
+
+            tvm_output = run_tvm_graph(tflite_model_buffer, tvm_in_data, in_node, target=device)
+            for i in range(len(tflite_output)):
+                if output_need_transpose_nchw:
+                    tvm.testing.assert_allclose(tflite_output[i],
+                                                np.transpose(tvm_output[i], axes=(0, 2, 3, 1)),
+                                                atol=1e-5, rtol=1e-5)
+                else:
+                    tvm.testing.assert_allclose(tflite_output[i], tvm_output[i],
+                                                atol=1e-5, rtol=1e-5)
+
+        sess.close()
+
+
+#######################################################################
+# Pooling
+# -------
+def _test_pooling_iteration(input_shape, **kwargs):
+    """ One iteration of pool operation with given shapes and attributes """
+
+    x = -np.arange(
+        np.prod(input_shape), dtype=np.float32).reshape(input_shape) - 1
+    tvm_data = np.transpose(x, axes=(0, 3, 1, 2))
+
+    with tf.Graph().as_default():
+        in_data = array_ops.placeholder(shape=input_shape, dtype='float32')
+        out = nn_ops.pool(in_data, **kwargs)
+
+        compare_tflite_with_tvm(x, tvm_data, 'Placeholder:0', [in_data], [out],
+                                output_need_transpose_nchw=True)
+
+
+def _test_pooling(input_shape, **kwargs):
+    _test_pooling_iteration(input_shape, **kwargs)
+
+
+def test_forward_pooling():
+    """ Pooling """
+
+    for pool_type in ['AVG', 'MAX']:
+        _test_pooling(input_shape=[2, 9, 10, 2],
+                      window_shape=[1, 1],
+                      padding='SAME',
+                      pooling_type=pool_type,
+                      dilation_rate=[1, 1],
+                      strides=[1, 1])
+
+        _test_pooling(input_shape=[2, 10, 9, 2],
+                      window_shape=[1, 1],
+                      padding='SAME',
+                      pooling_type=pool_type,
+                      dilation_rate=[1, 1],
+                      strides=[1, 1])
+
+        _test_pooling(input_shape=[2, 9, 10, 2],
+                      window_shape=[2, 1],
+                      padding='SAME',
+                      pooling_type=pool_type,
+                      dilation_rate=[1, 1],
+                      strides=[1, 1])
+
+        _test_pooling(input_shape=[2, 10, 9, 2],
+                      window_shape=[2, 3],
+                      padding='SAME',
+                      pooling_type=pool_type,
+                      dilation_rate=[1, 1],
+                      strides=[2, 1])
+
+
+#######################################################################
+# Convolution
+# -----------
+
+def _test_convolution(tensor_in_sizes, filter_in_sizes,
+                      dilations, strides, padding, data_format,
+                      is_depthwise=False):
+    """ One iteration of convolution with given shapes and attributes """
+
+    total_size_1 = 1
+    total_size_2 = 1
+    for s in tensor_in_sizes:
+        total_size_1 *= s
+    for s in filter_in_sizes:
+        total_size_2 *= s
+    # Initializes the input tensor with array containing incrementing
+    # numbers from 1.
+    data_array = [f * 1.0 for f in range(1, total_size_1 + 1)]
+    filter_array = [f * 1.0 for f in range(1, total_size_2 + 1)]
+
+    with tf.Graph().as_default():
+        in_data = array_ops.placeholder(shape=tensor_in_sizes, dtype='float32')
+        in_filter = constant_op.constant(filter_array, shape=filter_in_sizes, dtype='float32')
+        strides = [1] + strides + [1]
+        dilations = [1] + dilations + [1]
+
+        if is_depthwise:
+            out = nn_ops.depthwise_conv2d_native(in_data,
+                                                 in_filter,
+                                                 strides=strides,
+                                                 padding=padding,
+                                                 data_format=data_format)
+        else:
+            out = nn_ops.conv2d(in_data,
+                                in_filter,
+                                strides=strides,
+                                padding=padding,
+                                data_format=data_format)
+        # TFLite is NHWC, TVM is NCHW
+        tflite_data_array = np.reshape(data_array, tensor_in_sizes).astype('float32')
+        tvm_data_array = np.transpose(tflite_data_array, axes=(0, 3, 1, 2))
+        # TFLite output is NHWC, TVM is NCHW, we need transpose
+        compare_tflite_with_tvm(tflite_data_array, tvm_data_array,
+                                'Placeholder:0', [in_data], [out],
+                                output_need_transpose_nchw=True)
+
+
+def test_forward_convolution():
+    _test_convolution([4, 8, 8, 176], [1, 1, 176, 32], [1, 1], [1, 1], 'SAME', 'NHWC')
+    _test_convolution([4, 17, 17, 19], [3, 3, 19, 19], [1, 1], [2, 2], 'VALID', 'NHWC')
+    _test_convolution([4, 17, 17, 124], [1, 1, 124, 19], [1, 1], [1, 1], 'SAME', 'NHWC')
+    _test_convolution([4, 17, 17, 12], [3, 3, 12, 32], [1, 1], [2, 2], 'VALID', 'NHWC')
+
+    # depthwise convolution
+    _test_convolution([4, 8, 8, 176], [1, 1, 176, 1], [1, 1], [1, 1], 'SAME', 'NHWC', True)
+    _test_convolution([4, 17, 17, 19], [3, 3, 19, 1], [1, 1], [2, 2], 'VALID', 'NHWC', True)
+    _test_convolution([4, 17, 17, 124], [1, 1, 124, 1], [1, 1], [1, 1], 'SAME', 'NHWC', True)
+    _test_convolution([4, 17, 17, 12], [3, 3, 12, 1], [1, 1], [2, 2], 'VALID', 'NHWC', True)
+
+
+#######################################################################
+# Reshape
+# -------
+
+def _test_reshape(data, out_shape):
+    """ One iteration of reshape operation with given data and out shape """
+    # see relay/frontend/tflite.py convert_reshape more detail of channel first rule
+    if len(data.shape) == 1 or len(data.shape) == 2:
+        tvm_data = data
+    elif len(data.shape) == 3:
+        tvm_data = np.transpose(data, axes=(0, 2, 1))
+    elif len(data.shape) == 4:
+        tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
+    else:
+        raise NotImplementedError("Not support input shape {} of reshape : ".
+                                  format(str(len(data))))
+
+    with tf.Graph().as_default():
+        in_data = array_ops.placeholder(shape=data.shape, dtype=data.dtype)
+        out = array_ops.reshape(in_data, out_shape)
+
+        compare_tflite_with_tvm(data, tvm_data, 'Placeholder:0', [in_data], [out])
+
+
+def test_forward_reshape():
+    _test_reshape(np.arange(6.0, dtype=np.float32), [2, 3])
+    _test_reshape(np.arange(6), [-1, 2])
+    _test_reshape(np.arange(6), [3, -1])
+    _test_reshape(np.arange(6), [-1])
+
+
+#######################################################################
+# Squeeze
+# -------
+
+def _test_squeeze(data, squeeze_dims=None):
+    """ One iteration of squeeze """
+
+    if squeeze_dims is None:
+        squeeze_dims = []
+
+    # see relay/frontend/tflite.py convert_squeeze more detail of channel first rule
+    if len(data.shape) == 1 or len(data.shape) == 2:
+        tvm_data = data
+    elif len(data.shape) == 3:
+        tvm_data = np.transpose(data, axes=(0, 2, 1))
+    elif len(data.shape) == 4:
+        tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
+    else:
+        raise NotImplementedError("Not support input shape {} of reshape : ".
+                                  format(str(len(data.shape))))
+
+    tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
+
+    with tf.Graph().as_default():
+        in_data = array_ops.placeholder(shape=data.shape, dtype=data.dtype)
+
+        if squeeze_dims:
+            out = array_ops.squeeze(in_data, squeeze_dims)
+        else:
+            out = array_ops.squeeze(in_data)
+
+        compare_tflite_with_tvm(data, tvm_data, 'Placeholder:0', [in_data], [out])
+
+
+def test_forward_squeeze():
+    """ Squeeze """
+    _test_squeeze(np.arange(6).reshape((1, 2, 1, 3)), [0, 2])
+    _test_squeeze(np.arange(6).reshape((2, 1, 3, 1)), [1, 3])
+
+#######################################################################
+# Softmax
+# -------
+
+def _test_softmax(data):
+    """ One iteration of softmax """
+    with tf.Graph().as_default():
+        in_data = array_ops.placeholder(shape=data.shape, dtype=data.dtype)
+        out = nn_ops.softmax(in_data)
+        compare_tflite_with_tvm(data, data, 'Placeholder:0', [in_data], [out])
+
+def test_forward_softmax():
+    """ Softmax """
+    _test_softmax(np.arange(6.0, dtype=np.float32).reshape((1, 6)))
+
+#######################################################################
+# Mobilenet
+# ---------
+def test_forward_mobilenet():
+    '''test mobilenet v1 tflite model'''
+    # MobilenetV1
+    temp = util.tempdir()
+    tflite_model_file = nnvm.testing.tf.get_workload_official(
+        "http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224.tgz",
+        "mobilenet_v1_1.0_224.tflite", temp)
+    tflite_model_buf = open(tflite_model_file, "rb").read()
+    data = np.random.uniform(size=(1, 224, 224, 3)).astype('float32')
+    tvm_data = np.transpose(data, axes=(0, 3, 1, 2))
+    tflite_output = run_tflite_graph(tflite_model_buf, data)
+    tvm_output = run_tvm_graph(tflite_model_buf, tvm_data, 'input')
+    tvm.testing.assert_allclose(np.squeeze(tvm_output[0]), np.squeeze(tflite_output[0]),
+                                rtol=1e-5, atol=1e-5)
+
+#######################################################################
+# Main
+# ----
+if __name__ == '__main__':
+    # Transforms
+    test_forward_reshape()
+    test_forward_squeeze()
+
+    # NN
+    test_forward_convolution()
+    test_forward_pooling()
+    test_forward_softmax()
+
+    # End to End
+    test_forward_mobilenet()

tests/scripts/task_python_frontend.sh

@@ -38,3 +38,6 @@ python3 -m nose -v tests/python/frontend/onnx || exit -1
 
 echo "Running nnvm to relay frontend test..."
 python3 -m nose -v tests/python/frontend/nnvm_to_relay || exit -1
+
+echo "Running relay TFLite frontend test..."
+python3 -m nose -v tests/python/frontend/tflite || exit -1