from_tflite.py 6.73 KB
Newer Older
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.
17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36
"""
Compile TFLite Models
===================
**Author**: `Zhao Wu <https://github.com/FrozenGene>`_

This article is an introductory tutorial to deploy TFLite models with Relay.

To get started, Flatbuffers and TFLite package needs to be installed as prerequisites.

A quick solution is to install Flatbuffers via pip

.. code-block:: bash

    pip install flatbuffers --user

To install TFlite packages, you could use our prebuilt wheel:

.. code-block:: bash

    # For python3:
37
    wget https://raw.githubusercontent.com/dmlc/web-data/master/tensorflow/tflite/whl/tflite-0.0.1-py3-none-any.whl
38 39 40
    pip install tflite-0.0.1-py3-none-any.whl --user

    # For python2:
41
    wget https://raw.githubusercontent.com/dmlc/web-data/master/tensorflow/tflite/whl/tflite-0.0.1-py2-none-any.whl
42 43 44 45 46 47 48 49 50 51 52 53 54
    pip install tflite-0.0.1-py2-none-any.whl --user


or you could generate TFLite package by yourself. The steps are as following:

.. code-block:: bash

    # Get the flatc compiler.
    # Please refer to https://github.com/google/flatbuffers for details
    # and make sure it is properly installed.
    flatc --version

    # Get the TFLite schema.
55
    wget https://raw.githubusercontent.com/tensorflow/tensorflow/r1.13/tensorflow/lite/schema/schema.fbs
56 57 58 59 60 61 62 63 64 65 66 67 68 69 70

    # Generate TFLite package.
    flatc --python schema.fbs

    # Add it to PYTHONPATH.
    export PYTHONPATH=/path/to/tflite


Now please check if TFLite package is installed successfully, ``python -c "import tflite"``

Below you can find an example for how to compile TFLite model using TVM.
"""
######################################################################
# Utils for downloading and extracting zip files
# ---------------------------------------------
71
import os
72 73 74 75

def extract(path):
    import tarfile
    if path.endswith("tgz") or path.endswith("gz"):
76
        dir_path = os.path.dirname(path)
77
        tar = tarfile.open(path)
78
        tar.extractall(path=dir_path)
79 80 81 82 83 84 85 86 87
        tar.close()
    else:
        raise RuntimeError('Could not decompress the file: ' + path)


######################################################################
# Load pretrained TFLite model
# ---------------------------------------------
# we load mobilenet V1 TFLite model provided by Google
88 89
from tvm.contrib.download import download_testdata

90 91 92
model_url = "http://download.tensorflow.org/models/mobilenet_v1_2018_08_02/mobilenet_v1_1.0_224.tgz"

# we download model tar file and extract, finally get mobilenet_v1_1.0_224.tflite
93 94 95
model_path = download_testdata(model_url, "mobilenet_v1_1.0_224.tgz", module=['tf', 'official'])
model_dir = os.path.dirname(model_path)
extract(model_path)
96 97

# now we have mobilenet_v1_1.0_224.tflite on disk and open it
98
tflite_model_file = os.path.join(model_dir, "mobilenet_v1_1.0_224.tflite")
99 100 101 102 103 104 105 106 107 108 109 110 111 112 113
tflite_model_buf = open(tflite_model_file, "rb").read()

# get TFLite model from buffer
import tflite.Model
tflite_model = tflite.Model.Model.GetRootAsModel(tflite_model_buf, 0)

######################################################################
# Load a test image
# ---------------------------------------------
# A single cat dominates the examples!
from PIL import Image
from matplotlib import pyplot as plt
import numpy as np

image_url = 'https://github.com/dmlc/mxnet.js/blob/master/data/cat.png?raw=true'
114 115
image_path = download_testdata(image_url, 'cat.png', module='data')
resized_image = Image.open(image_path).resize((224, 224))
116 117 118 119
plt.imshow(resized_image)
plt.show()
image_data = np.asarray(resized_image).astype("float32")

120
# after expand_dims, we have format NHWC
121 122 123 124
image_data = np.expand_dims(image_data, axis=0)

# preprocess image as described here:
# https://github.com/tensorflow/models/blob/edb6ed22a801665946c63d650ab9a0b23d98e1b1/research/slim/preprocessing/inception_preprocessing.py#L243
125 126 127
image_data[:, :, :, 0] = 2.0 / 255.0 * image_data[:, :, :, 0] - 1
image_data[:, :, :, 1] = 2.0 / 255.0 * image_data[:, :, :, 1] - 1
image_data[:, :, :, 2] = 2.0 / 255.0 * image_data[:, :, :, 2] - 1
128 129 130 131 132 133 134 135
print('input', image_data.shape)

######################################################################
# Compile the model with relay
# ---------------------------------------------

# TFLite input tensor name, shape and type
input_tensor = "input"
136
input_shape = (1, 224, 224, 3)
137 138 139 140
input_dtype = "float32"

# parse TFLite model and convert into Relay computation graph
from tvm import relay
141 142 143
mod, params = relay.frontend.from_tflite(tflite_model,
                                         shape_dict={input_tensor: input_shape},
                                         dtype_dict={input_tensor: input_dtype})
144

145
# target x86 CPU
146
target = "llvm"
147
with relay.build_config(opt_level=3):
148
    graph, lib, params = relay.build(mod[mod.entry_func], target, params=params)
149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180

######################################################################
# Execute on TVM
# ---------------------------------------------
import tvm
from tvm.contrib import graph_runtime as runtime

# create a runtime executor module
module = runtime.create(graph, lib, tvm.cpu())

# feed input data
module.set_input(input_tensor, tvm.nd.array(image_data))

# feed related params
module.set_input(**params)

# run
module.run()

# get output
tvm_output = module.get_output(0).asnumpy()

######################################################################
# Display results
# ---------------------------------------------

# load label file
label_file_url = ''.join(['https://raw.githubusercontent.com/',
                          'tensorflow/tensorflow/master/tensorflow/lite/java/demo/',
                          'app/src/main/assets/',
                          'labels_mobilenet_quant_v1_224.txt'])
label_file = "labels_mobilenet_quant_v1_224.txt"
181
label_path = download_testdata(label_file_url, label_file, module='data')
182

183
# list of 1001 classes
184
with open(label_path) as f:
185
    labels = f.readlines()
186 187 188 189 190 191 192 193 194

# convert result to 1D data
predictions = np.squeeze(tvm_output)

# get top 1 prediction
prediction = np.argmax(predictions)

# convert id to class name and show the result
print("The image prediction result is: id " + str(prediction) + " name: " + labels[prediction])