In [0]:
!pip install tf-nightly
In [0]:
# This is important!
import os
os.environ['TF_ENABLE_CONTROL_FLOW_V2'] = '1'
import tensorflow as tf
import numpy as np
Note we will be using tf.lite.experimental.nn.TFLiteLSTMCell
& tf.lite.experimental.nn.dynamic_rnn
in the tutorial.
Also note here, we're not trying to build the model to be a real world application, but only demonstrates how to use TensorFlow lite. You can a build a much better model using CNN models.
For more canonical lstm codelab, please see here.
In [0]:
# Step 1: Build the MNIST LSTM model.
def buildLstmLayer(inputs, num_layers, num_units):
"""Build the lstm layer.
Args:
inputs: The input data.
num_layers: How many LSTM layers do we want.
num_units: The unmber of hidden units in the LSTM cell.
"""
lstm_cells = []
for i in range(num_layers):
lstm_cells.append(
tf.lite.experimental.nn.TFLiteLSTMCell(
num_units, forget_bias=0, name='rnn{}'.format(i)))
lstm_layers = tf.keras.layers.StackedRNNCells(lstm_cells)
# Assume the input is sized as [batch, time, input_size], then we're going
# to transpose to be time-majored.
transposed_inputs = tf.transpose(
inputs, perm=[1, 0, 2])
outputs, _ = tf.lite.experimental.nn.dynamic_rnn(
lstm_layers,
transposed_inputs,
dtype='float32',
time_major=True)
unstacked_outputs = tf.unstack(outputs, axis=0)
return unstacked_outputs[-1]
tf.reset_default_graph()
model = tf.keras.models.Sequential([
tf.keras.layers.Input(shape=(28, 28), name='input'),
tf.keras.layers.Lambda(buildLstmLayer, arguments={'num_layers' : 2, 'num_units' : 64}),
tf.keras.layers.Flatten(),
tf.keras.layers.Dense(10, activation=tf.nn.softmax, name='output')
])
model.compile(optimizer='adam',
loss='sparse_categorical_crossentropy',
metrics=['accuracy'])
model.summary()
In [0]:
# Step 2: Train & Evaluate the model.
(x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0
# Cast x_train & x_test to float32.
x_train = x_train.astype(np.float32)
x_test = x_test.astype(np.float32)
model.fit(x_train, y_train, epochs=5)
model.evaluate(x_test, y_test)
In [0]:
# Step 3: Convert the Keras model to TensorFlow Lite model.
sess = tf.keras.backend.get_session()
input_tensor = sess.graph.get_tensor_by_name('input:0')
output_tensor = sess.graph.get_tensor_by_name('output/Softmax:0')
converter = tf.lite.TFLiteConverter.from_session(
sess, [input_tensor], [output_tensor])
tflite = converter.convert()
print('Model converted successfully!')
In [0]:
# Step 4: Check the converted TensorFlow Lite model.
interpreter = tf.lite.Interpreter(model_content=tflite)
try:
interpreter.allocate_tensors()
except ValueError:
assert False
MINI_BATCH_SIZE = 1
correct_case = 0
for i in range(len(x_test)):
input_index = (interpreter.get_input_details()[0]['index'])
interpreter.set_tensor(input_index, x_test[i * MINI_BATCH_SIZE: (i + 1) * MINI_BATCH_SIZE])
interpreter.invoke()
output_index = (interpreter.get_output_details()[0]['index'])
result = interpreter.get_tensor(output_index)
# Reset all variables so it will not pollute other inferences.
interpreter.reset_all_variables()
# Evaluate.
prediction = np.argmax(result)
if prediction == y_test[i]:
correct_case += 1
print('TensorFlow Lite Evaluation result is {}'.format(correct_case * 1.0 / len(x_test)))