-
Notifications
You must be signed in to change notification settings - Fork 10
TensorFlow Lite Backend
These instructions are for TF 1.15. Instructions for TF 2.0 and higher are forthcoming.
TFLite is a lightweight library used to run a subset of TensorFlow models in constrained environments, such as mobile devices. Tensor/IO includes support for TF Lite models, but to use your models with the TF Lite backend you must first convert them to that format.
This process assumes you have already converted your model to the saved graph format from training checkpoints using tf.estimator.export_saved_model or an equivalent function. For more information see TF Estimator export_saved_model.
First make sure you have the tensorflow python package installed, and at the command line, confirm that the saved_model_cli, freeze_graph, and toco commands are available:
$ which saved_model_cli
/path/to/tensorflow/bin/saved_model_cli
$ which freeze_graph
/path/to/tensorflow/bin/freeze_graph
# which toco
/path/to/tensorflow/bin/toco
We'll be using environmental variables in the script to improve readability. Set MODEL_DIR to the saved model's directory:
$ MODEL_DIR=/path/to/model
The first step is to get information about your model's input and output layers using tensorflow's saved model cli. We need the name and the shape of our input and output nodes for the next steps.
$ saved_model_cli show --all --dir $MODEL_DIR
MetaGraphDef with tag-set: 'serve' contains the following SignatureDefs:
signature_def['serving_default']:
The given SavedModel SignatureDef contains the following input(s):
inputs['image'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 128, 128, 3)
name: image:0
The given SavedModel SignatureDef contains the following output(s):
outputs['class'] tensor_info:
dtype: DT_INT32
shape: (-1, 1)
name: ToInt32_1:0
outputs['probability'] tensor_info:
dtype: DT_FLOAT
shape: (-1, 1)
name: sigmoid:0
In my case I have an input layer named image (the trailing ':0' can be disregarded unless you have more than one layer with this name), and two output layers. The one I'm interested in is sigmoid. Disregarding the batch dimension, the shapes are [128,128,3] and [1], respectively, although we won't use the output shape in the next command.
If you receive a "Shapes must be equal rank" error when running the toco command below, trying including a batch dimension with a size of one: [1,128,128,3].
Set the following environmental variables to those values:
$ INPUT_NODES=image
$ OUTPUT_NODES=sigmoid
$ INPUT_SHAPES=128,128,3
The next step is to freeze the graph. Freezing the graph converts the model variables into constants and embeds them into a single file along with the graph definition.
First set an environmental variable for the output of this command:
$ OUTPUT_GRAPH=frozen_graph.pb
And then freeze it:
$ freeze_graph --input_saved_model_dir $MODEL_DIR --output_graph $OUTPUT_GRAPH --output_node_names $OUTPUT_NODES
Finally we'll use the toco command line utility to convert our frozen graph into a tensorflow lite compatible file. We'll use the environmental variables we set up earlier, in addition to one for the output of this command:
$ OUTPUT_FILE=model.tflite
And perform the conversion. I've broken this command onto multiple lines:
$ toco \
--graph_def_file $OUTPUT_GRAPH \
--output_file $OUTPUT_FILE \
--inference_type FLOAT \
--input_arrays $INPUT_NODES \
--input_shapes $INPUT_SHAPES \
--output_arrays $OUTPUT_NODES
And that's it. You should end up with a file in your working directory with whatever name you set OUTPUT_FILE to, in our case, model.tflite.
And with that file, you're ready to package your model.
Last updated 2/5/20.