Starting with a simple model:

As a prerequisite, I wanted to choose a TensorFlow model that wasn’t pre-trained or converted into a .tflite file already, so naturally I landed on a simple neural network trained on MNIST data (currently there are 3 TensorFlow Lite models supported: MobileNet, Inception v3, and On Device Smart Reply).



Luckily, Google open sources plenty of research and available models in its model zoo, including a MNIST training script. We’ll be referencing that code in this section so take a quick look to familiarize yourself with it.

There are a couple of modifications we should make to this training script in order to make the mobile conversion easier later on.

1. The first question to ask is “What are the input and output layers?”

From this, we clearly see that the input layer is tf.reshape so let’s give it a name.

y = tf.reshape(inputs, self._input_shape, name='input_tensor’)

A good practice is to name your input and output layers. This will save you some time and energy later on so you won’t have to dig around in TensorBoard in order to fill out certain parameters for the conversion tools. (an added benefit is that if you share your model without the training script, developers can examine it and quickly identify the I/O of your graph).

2. For the output layer, my first guess was the logit layer, but that isn’t the one we want to capture for inference. Instead, we want the softmax layer for the output of our inference model. (see below on line 7)

We also need to expose the softmax layer in the TensorFlow Graph because it’s the output layer used for inference. Right now it’s embedded inside the inference method. As a simple fix, move it out so when we train this model, the graph will include this layer. There are obviously more elegant ways to change this, but this is a simple way to edit the existing MNIST script.

To summarize, we looked at the training script and specifically named our input and output layers that we needed for model inference. It’s important to know the difference between training and inference layers. Since we’d like our prepared model to only be used for inference on mobile (predicting handwritten numbers in the case of MNIST data), we only need the layers necessary for that prediction.

Keep in mind that the MNIST script that we’re using does both (the training and inference). We’ll see the separation between both in TensorBoard a little later.

Here’s the completely edited mnist.py file for your reference.

To train the model, run the following (from the model project root). This takes about 1–2 hours with my 2017 Macbook Pro.

python official/mnist/mnist.py --export_dir /tmp/mnist_saved_model --model-dir /tmp/mnist_graph_def_with_ckpts

These export directories save checkpoints and protobuf files that define the graph. Let’s dissect the different TensorFlow formats that are saved from the training script.