Convert TensorFlow 2 Model to TensorFlow Lite Model

TensorFlow 2 provides the TFLiteConverter which allows converting a TensorFlow 2 model to TensorFlow Lite model. This model uses the .tflite file extension. TensorFlow Lite models can be executed using TensorFlow Lite interpreter without installing all TensorFlow packages.

Before starting, make sure you have installed tensorflow using pip package manager.

pip install tensorflow

We have created a simple model which consists of one layer which has one neuron. This model allows solving simple linear regression problem. It predicts the value of y for the given value of x. Relationship between x and y can be described as y = 2 * x + 1. A trained model is saved in HDF5 format.

train.py

from tensorflow import keras
import numpy as np

xs = np.array([-2.0, -1.0, 0.0, 1.0, 2.0, 3.0, 4.0], dtype=float)
ys = np.array([-3.0, -1.0, 1.0, 3.0, 5.0, 7.0, 9.0], dtype=float)

model = keras.Sequential([
    keras.layers.Dense(units=1, input_shape=[1])
])

model.compile(optimizer='sgd', loss='mean_squared_error')

model.fit(xs, ys, epochs=400)

model.save('model.h5')

We load the model from HDF5 format file. Then the model is converted to TensorFlow Lite model and saved to a .tflite file.

convert.py

import tensorflow as tf

model = tf.keras.models.load_model('model.h5')
converter = tf.lite.TFLiteConverter.from_keras_model(model)

tfliteModel = converter.convert()

with open('model.tflite', 'wb') as f:
    f.write(tfliteModel)

Now we can use the TensorFlow Lite model for predicting value of y for a previously unknown value of x. A model is loaded into memory using tf.lite.Interpreter. We set a value of the input tensor. A model is executed and the value of the output tensor is printed.

test.py

import tensorflow as tf
import numpy as np

interpreter = tf.lite.Interpreter('model.tflite')
interpreter.allocate_tensors()

x = np.float32([[15.0]])

inputDetails = interpreter.get_input_details()
interpreter.set_tensor(inputDetails[0]['index'], x)

interpreter.invoke()

outputDetails = interpreter.get_output_details()
y = interpreter.get_tensor(outputDetails[0]['index'])[0]

print(y[0])

In this case, x is equal to 15.0. Our model returns that y is 31.001108. We can verify result:

y = 2 * x + 1 = 2 * 15 + 1 = 31