Step 1. Save the Model

The TensorFlow Lite converter works on a number of different file formats, including SavedModel (preferred) and the Keras H5 format. For this exercise we’ll use SavedModel.

To do this, simply specify a directory to save the model in and call tf.saved_model.save, passing it the model and directory:

	export_dir = 'saved_model/1'
tf.saved_model.save(model, export_dir)

The model will be saved out as assets and variables as well as a saved_model.pb file, as shown in Figure 12-2.

Figure 12-2. SavedModel structure

Once you have the saved model, you can convert it using the TensorFlow Lite converter.

Step 2. Convert and Save the Model

The TensorFlow Lite converter is in the tf.lite package. You can call it to convert a saved model by first invoking it with the from_saved_model method, passing it the directory containing the saved model, and then invoking its convert method:

	# Convert the model.
converter = tf.lite.TFLiteConverter.from_saved_model(export_dir)
tflite_model = converter.convert()

You can then save out the new .tflite model using pathlib:

	import pathlib
tflite_model_file = pathlib.Path('model.tflite')
tflite_model_file.write_bytes(tflite_model)

At this point you have a .tflite file that you can use in any of the interpreter environments. Later we’ll use it on Android and iOS, but for now, let’s use the Python-based interpreter so you can run it in Colab. This same interpreter can be used in embedded Linux environments like a Raspberry Pi!

Step 3. Load the TFLite Model and Allocate Tensors

The next step is to load the model into the interpreter, allocate tensors that will be used for inputting data to the model for prediction, and then read the predictions that the model outputs. This is where using TensorFLow Lite, from a programmer’s perspective, greatly differs from using TensorFlow. With TensorFlow you can just say model.predict(something) and get the results, but because TensorFlow Lite won’t have many of the dependencies that TensorFlow does, particularly in non-Python environments, you now have to get a bit more low-level and deal with the input and output tensors, formatting your data to fit them and parsing the output in a way that makes sense for your device.

First, load the model and allocate the tensors:

	interpreter = tf.lite.Interpreter(model_content=tflite_model)
interpreter.allocate_tensors()

Then you can get the input and output details from the model, so you can begin to understand what data format it expects and what data format it will provide back to you:

	input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()
print(input_details)
print(output_details)

You’ll get a lot of output!

First, let’s inspect the input parameter. Note the shape setting, which is an array of type [1,1]. Also note the class, which is numpy.float32. These settings will dictate the shape of the input data and its format:

	[{'name': 'dense_input', 'index': 0, 'shape': array([1, 1], dtype=int32), 
  'shape_signature': array([1, 1], dtype=int32), 'dtype': <class 
  'numpy.float32'>, 'quantization': (0.0, 0), 'quantization_parameters': 
  {'scales': array([], dtype=float32), 'zero_points': array([], dtype=int32), 
  'quantized_dimension': 0}, 'sparsity_parameters': {}}]

So, in order to format the input data, you’ll need to use code like this to define the input array shape and type if you want to predict the y for x=10.0:

	to_predict = np.array([[10.0]], dtype=np.float32)
print(to_predict)

The double brackets around the 10.0 can cause a little confusion—the mnemonic I use for the array[1,1] here is to say that there is 1 list, giving us the first set of [], and that list contains just 1 value, which is [10.0], thus giving [[10.0]]. It can also be confusing that the shape is defined as dtype=int32, whereas you’re using numpy.float32. The dtype parameter is the data type defining the shape, not the contents of the list that is encapsulated in that shape. For that, you’ll use the class.

The output details are very similar, and what you want to keep an eye on here is the shape. As it’s also an array of type [1,1], you can expect the answer to be [[y]] in much the same way as the input was [[x]]:

	[{'name': 'Identity', 'index': 3, 'shape': array([1, 1], dtype=int32), 
  'shape_signature': array([1, 1], dtype=int32), 'dtype': <class 
  'numpy.float32'>, 'quantization': (0.0, 0), 'quantization_parameters': 
  {'scales': array([], dtype=float32), 'zero_points': array([], dtype=int32), 
  'quantized_dimension': 0}, 'sparsity_parameters': {}}]

Step 4. Perform the Prediction

To get the interpreter to do the prediction, you set the input tensor with the value to predict, telling it what input value to use:

	interpreter.set_tensor(input_details[0]['index'], to_predict)
interpreter.invoke()

The input tensor is specified using the index of the array of input details. In this case you have a very simple model that has only a single input option, so it’s input_details[0], and you’ll address it at the index. Input details item 0 has only one index, indexed at 0, and it expects a shape of [1,1] as defined earlier. So, you put the to_predict value in there. Then you invoke the interpreter with the invoke method.

You can then read the prediction by calling get_tensor and supplying it with the details of the tensor you want to read:

	tflite_results = interpreter.get_tensor(output_details[0]['index'])
print(tflite_results)

Again, there’s only one output tensor, so it will be output_details[0], and you specify the index to get the details beneath it, which will have the output value.

So, for example, if you run this code:

	to_predict = np.array([[10.0]], dtype=np.float32)
print(to_predict)
interpreter.set_tensor(input_details[0]['index'], to_predict)
interpreter.invoke()
tflite_results = interpreter.get_tensor(output_details[0]['index'])
print(tflite_results)

you should see output like:

	[[10.]]
[[18.975412]]

where 10 is the input value and 18.97 is the predicted value, which is very close to 19, which is 2X – 1 when X = 10. For why it’s not 19, look back to Chapter 1!

Given that this is a very simple example, let’s look at something a little more complex next—using transfer learning on a well-known image classification model, and then converting that for TensorFlow Lite. From there we’ll also be able to better explore the impacts of optimizing and quantizing the model.

Step 1. Build and Save the Model

First, get all of the data:

	import numpy as np
import matplotlib.pylab as plt

import tensorflow as tf
import tensorflow_hub as hub
import tensorflow_datasets as tfds


def format_image(image, label):
    image = tf.image.resize(image, IMAGE_SIZE) / 255.0
    return  image, label


(raw_train, raw_validation, raw_test), metadata = tfds.load(
    'cats_vs_dogs',
    split=['train[:80%]', 'train[80%:90%]', 'train[90%:]'],
    with_info=True,
    as_supervised=True,
)

num_examples = metadata.splits['train'].num_examples
num_classes = metadata.features['label'].num_classes
print(num_examples)
print(num_classes)

BATCH_SIZE = 32
train_batches = 
raw_train.shuffle(num_examples // 4)
.map(format_image).batch(BATCH_SIZE).prefetch(1)

validation_batches = raw_validation.map(format_image)
.batch(BATCH_SIZE).prefetch(1)
test_batches = raw_test.map(format_image).batch(1)

This will download the Dogs vs. Cats dataset and split it into training, test, and validation sets.

Next, you’ll use the mobilenet_v2 model from TensorFlow Hub to create a Keras layer called feature_extractor:

	module_selection = ("mobilenet_v2", 224, 1280) 
handle_base, pixels, FV_SIZE = module_selection

MODULE_HANDLE ="https://tfhub.dev/google/tf2-preview/{}/feature_vector/4"
.format(handle_base)

IMAGE_SIZE = (pixels, pixels)

feature_extractor = hub.KerasLayer(MODULE_HANDLE,
                                   input_shape=IMAGE_SIZE + (3,), 
                                   output_shape=[FV_SIZE],
                                   trainable=False)

Now that you have the feature extractor, you can make it the first layer in your neural network and add an output layer with as many neurons as you have classes (in this case, two). You can then compile and train it:

	model = tf.keras.Sequential([
            feature_extractor,
            tf.keras.layers.Dense(num_classes, activation='softmax')
        ])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

hist = model.fit(train_batches,
                 epochs=5,
                 validation_data=validation_batches)

With just five epochs of training, this should give a model with 99% accuracy on the training set and 98%+ on the validation set. Now you can simply save the model out:

	CATS_VS_DOGS_SAVED_MODEL = "exp_saved_model"
tf.saved_model.save(model, CATS_VS_DOGS_SAVED_MODEL)

Once you have the saved model, you can convert it.

Step 2. Convert the Model to TensorFlow Lite

As before, you can now take the saved model and convert it into a .tflite model. You’ll save it out as converted_model.tflite:

	converter = tf.lite.TFLiteConverter.from_saved_model(CATS_VS_DOGS_SAVED_MODEL)
tflite_model = converter.convert()
tflite_model_file = 'converted_model.tflite'

with open(tflite_model_file, "wb") as f:
    f.write(tflite_model)

Once you have the file, you can instantiate an interpreter with it. When this is done, you should get the input and output details as before. Load them into variables called input_index and output_index, respectively. This makes the code a little more readable!

	interpreter = tf.lite.Interpreter(model_path=tflite_model_file)
interpreter.allocate_tensors()

input_index = interpreter.get_input_details()[0]["index"]
output_index = interpreter.get_output_details()[0]["index"]

predictions = []

The dataset has lots of test images in test_batches, so if you want to take 100 of the images and test them, you can do so like this (feel free to change the 100 to any other value):

	test_labels, test_imgs = [], []
for img, label in test_batches.take(100):
    interpreter.set_tensor(input_index, img)
    interpreter.invoke()
    predictions.append(interpreter.get_tensor(output_index))
    test_labels.append(label.numpy()[0])
    test_imgs.append(img)

Earlier, when reading the images, they were reformatted by the mapping function called format_image to be the right size for both training and inference, so all you have to do now is set the interpreter’s tensor at the input index to the image. After invoking the interpreter, you can then get the tensor at the output index.

If you want to see how the predictions did against the labels, you can run code like this:

	score = 0
for item in range(0,99):
    prediction=np.argmax(predictions[item])
    label = test_labels[item]
    if prediction==label:
        score=score+1

print("Out of 100 predictions I got " + str(score) + " correct")

This should give you a score of 99 or 100 correct predictions.

You can also visualize the output of the model against the test data with this code:

	for index in range(0,99):
    plt.figure(figsize=(6,3))
    plt.subplot(1,2,1)
    plot_image(index, predictions, test_labels, test_imgs)
    plt.show()

You can see some of the results of this in Figure 12-3. (Note that all the code is available in the book’s GitHub repo, so if you need it, take a look there!)

Figure 12-3. Results of inference

This is just the plain, converted model without any optimizations for mobile added. In the next step you’ll explore how to optimize this model for mobile devices.

Step 3. Optimize the Model

Now that you’ve seen the end-to-end process of training, converting, and using a model with the TensorFlow Lite interpreter, let’s look at how to get started with optimizing and quantizing the model.

The first type of optimization, called dynamic range quantization, is achieved by setting the optimizations property on the converter prior to performing the conversion. Here’s the code:

	converter = tf.lite.TFLiteConverter.from_saved_model(CATS_VS_DOGS_SAVED_MODEL)
converter.optimizations = [tf.lite.Optimize.DEFAULT]

tflite_model = converter.convert()
tflite_model_file = 'converted_model.tflite'

with open(tflite_model_file, "wb") as f:
    f.write(tflite_model)

There are several optimization options available at the time of writing (more may be added later). These include:

OPTIMIZE_FOR_SIZE

Perform optimizations that make the model as small as possible.

OPTIMIZE_FOR_LATENCY

Perform optimizations that reduce inference time as much as possible.

DEFAULT

Find the best balance between size and latency.

In this case, the model size was close to 9 MB before this step but only 2.3 MB afterwards—a reduction of almost 70%. Various experiments have shown that models can be made up to 4× smaller, with a 2–3× speedup. Depending on the model type, however, there can be a loss in accuracy, so it’s a good idea to test the model thoroughly if you quantize like this. In this case, I found that the accuracy of the model dropped from 99% to about 94%.

You can enhance this with full integer quantization or float16 quantization to take advantage of specific hardware. Full integer quantization changes the weights in the model from 32-bit floating point to 8-bit integer, which (particularly for larger models) can have a huge impact on model size and latency with a relatively small impact on accuracy.

To get full integer quantization, you’ll need to specify a representative dataset that tells the convertor roughly what range of data to expect. Update the code as follows:

	converter = tf.lite.TFLiteConverter.from_saved_model(CATS_VS_DOGS_SAVED_MODEL)

converter.optimizations = [tf.lite.Optimize.DEFAULT]

def representative_data_gen():
    for input_value, _ in test_batches.take(100):
        yield [input_value]

converter.representative_dataset = representative_data_gen
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8]

tflite_model = converter.convert()
tflite_model_file = 'converted_model.tflite'

with open(tflite_model_file, "wb") as f:
    f.write(tflite_model)

Having this representative data allows the convertor to inspect data as it flows through the model and find where to best make the conversions. Then, by setting the supported ops (in this case to INT8), you can ensure that the precision of only those parts of the model is quantized. The result might be a slightly larger model—in this case, it went from 2.3 MB when using convertor.optimizations only to 2.8 MB. However, the accuracy went back up to 99%. Thus, by following these steps you can reduce the model’s size by about two-thirds, while maintaining its accuracy!