from_tflite.py 7.02 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 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70
    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.
    wget https://raw.githubusercontent.com/tensorflow/tensorflow/r1.12/tensorflow/contrib/lite/schema/schema.fbs

    # 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 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 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 181 182 183 184 185 186 187 188 189
plt.imshow(resized_image)
plt.show()
image_data = np.asarray(resized_image).astype("float32")

# convert HWC to CHW
image_data = image_data.transpose((2, 0, 1))

# after expand_dims, we have format NCHW
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
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
print('input', image_data.shape)

####################################################################
#
# .. note:: Input layout:
#
#   Currently, TVM TFLite frontend accepts ``NCHW`` as input layout.

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

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

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

# targt x86 cpu
target = "llvm"
with relay.build_module.build_config(opt_level=3):
    graph, lib, params = relay.build(func, target, params=params)

######################################################################
# 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"
190
label_path = download_testdata(label_file_url, label_file, module='data')
191 192 193

# map id to 1001 classes
labels = dict()
194
with open(label_path) as f:
195 196 197 198 199 200 201 202 203 204 205
    for id, line in enumerate(f):
        labels[id] = line

# 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])