Converting Keras RNN model to TensorFlow Lite model in TF2

Question

I am currently trying to convert a RNN model to TF lite. After multiple failed attempts I tried running the example given in the repository found here . This threw errors too due to changes in the layer definition location. Once fixed to the code below

import os
os.environ['TF_ENABLE_CONTROL_FLOW_V2'] = '1'

import tensorflow as tf
import numpy as np
from tensorflow_core.lite.experimental.examples.lstm.rnn_cell import TFLiteLSTMCell
from tensorflow_core.lite.experimental.examples.lstm.rnn import dynamic_rnn

# 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)))
        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(
  outputs, _ = 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()

# 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)


# 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!')

# 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)))

I continue to get errors

Traceback (most recent call last):
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 2346, in _enter_scope_uncached
    entered_pure_variable_scope = pure_variable_scope.__enter__()
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 1912, in __enter__
    constraint=self._constraint)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/variable_scope.py", line 1073, in __init__
    raise NotImplementedError("Caching devices is not yet supported "
NotImplementedError: Caching devices is not yet supported when eager execution is enabled.

and after disabling the eager execution using tf.compat.v1.disable_eager_execution() I get the error

Traceback (most recent call last):
  File "/home/xyz/genesis-dnn-se/TFLite_example.py", line 46, in <module>
    tf.keras.layers.Dense(10, activation=tf.nn.softmax, name='output')
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/sequential.py", line 114, in __init__
    self.add(layer)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/training/tracking/base.py", line 457, in _method_wrapper
    result = method(self, *args, **kwargs)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/sequential.py", line 196, in add
    output_tensor = layer(self.outputs[0])
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 847, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/core.py", line 795, in call
    return self.function(inputs, **arguments)
  File "/home/xyz/genesis-dnn-se/TFLite_example.py", line 37, in buildLstmLayer
    time_major=True)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/lite/experimental/examples/lstm/rnn.py", line 266, in dynamic_rnn
    dtype=dtype)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/rnn.py", line 916, in _dynamic_rnn_loop
    swap_memory=swap_memory)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/control_flow_ops.py", line 2675, in while_loop
    back_prop=back_prop)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/while_v2.py", line 198, in while_loop
    add_control_dependencies=add_control_dependencies)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/framework/func_graph.py", line 915, in func_graph_from_py_func
    func_outputs = python_func(*func_args, **func_kwargs)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/while_v2.py", line 176, in wrapped_body
    outputs = body(*_pack_sequence_as(orig_loop_vars, args))
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/rnn.py", line 884, in _time_step
    (output, new_state) = call_cell()
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/ops/rnn.py", line 870, in <lambda>
    call_cell = lambda: cell(input_t, state)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/engine/base_layer.py", line 847, in __call__
    outputs = call_fn(cast_inputs, *args, **kwargs)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/python/keras/layers/recurrent.py", line 137, in call
    inputs, states = cell.call(inputs, states, **kwargs)
  File "/home/xyz/anaconda3/envs/tf2/lib/python3.7/site-packages/tensorflow_core/lite/experimental/examples/lstm/rnn_cell.py", line 440, in call
    if input_size.value is None:
AttributeError: 'int' object has no attribute 'value'

Does anyone have a working example of converting a RNN (LSTM, GRU, CustomRNN) to TFLite in TensorFlow 2.0. And I am using TF version 2.0.0.

Answer 1

import os
os.environ['TF_ENABLE_CONTROL_FLOW_V2'] = '1'
import sys
import tensorflow as tf
import argparse
tf.compat.v1.disable_eager_execution()

class MnistLstmModel(object):
  """Build a simple LSTM based MNIST model.

  Attributes:
    time_steps: The maximum length of the time_steps, but since we're just using
      the 'width' dimension as time_steps, it's actually a fixed number.
    input_size: The LSTM layer input size.
    num_lstm_layer: Number of LSTM layers for the stacked LSTM cell case.
    num_lstm_units: Number of units in the LSTM cell.
    units: The units for the last layer.
    num_class: Number of classes to predict.
  """

  def __init__(self, time_steps, input_size, num_lstm_layer, num_lstm_units,
               units, num_class):
    self.time_steps = time_steps
    self.input_size = input_size
    self.num_lstm_layer = num_lstm_layer
    self.num_lstm_units = num_lstm_units
    self.units = units
    self.num_class = num_class

  def build_model(self):
    """Build the model using the given configs.

    Returns:
      x: The input placeholder tensor.
      logits: The logits of the output.
      output_class: The prediction.
    """
    x = tf.compat.v1.placeholder(
        'float32', [None, self.time_steps, self.input_size], name='INPUT')
    lstm_layers = []
    for _ in range(self.num_lstm_layer):
      lstm_layers.append(
          # Important:
          #
          # Note here, we use `tf.lite.experimental.nn.TFLiteLSTMCell`
          # (OpHinted LSTMCell).
          tf.compat.v1.lite.experimental.nn.TFLiteLSTMCell(
              self.num_lstm_units, forget_bias=0))
    # Weights and biases for output softmax layer.
    out_weights = tf.Variable(tf.random.normal([self.units, self.num_class]))
    out_bias = tf.Variable(tf.zeros([self.num_class]))

    # Transpose input x to make it time major.
    lstm_inputs = tf.transpose(x, perm=[1, 0, 2])
    lstm_cells = tf.keras.layers.StackedRNNCells(lstm_layers)
    # Important:
    #
    # Note here, we use `tf.lite.experimental.nn.dynamic_rnn` and `time_major`
    # is set to True.
    outputs, _ = tf.compat.v1.lite.experimental.nn.dynamic_rnn(
        lstm_cells, lstm_inputs, dtype='float32', time_major=True)

    # Transpose the outputs back to [batch, time, output]
    outputs = tf.transpose(outputs, perm=[1, 0, 2])
    outputs = tf.unstack(outputs, axis=1)
    logits = tf.matmul(outputs[-1], out_weights) + out_bias
    output_class = tf.nn.softmax(logits, name='OUTPUT_CLASS')

    return x, logits, output_class

def train(model,
      model_dir,
      batch_size=20,
      learning_rate=0.001,
      train_steps=200,
      eval_steps=500,
      save_every_n_steps=1000):
  """Train & save the MNIST recognition model."""
  # Train & test dataset.
  (x_train, y_train), (x_test, y_test) = tf.keras.datasets.mnist.load_data()
  train_dataset = tf.compat.v1.data.Dataset.from_tensor_slices((x_train, y_train))
  train_iterator = train_dataset.shuffle(
      buffer_size=1000).batch(batch_size).repeat().make_one_shot_iterator()
  x, logits, output_class = model.build_model()
  test_dataset = tf.compat.v1.data.Dataset.from_tensor_slices((x_test, y_test))
  test_iterator = test_dataset.batch(
      batch_size).repeat().make_one_shot_iterator()
  # input label placeholder
  y = tf.compat.v1.placeholder(tf.int32, [
      None,
  ])
  one_hot_labels = tf.one_hot(y, depth=model.num_class)
  # Loss function
  loss = tf.reduce_mean(
      tf.nn.softmax_cross_entropy_with_logits(
          logits=logits, labels=one_hot_labels))
  correct = tf.compat.v1.nn.in_top_k(output_class, y, 1)
  accuracy = tf.reduce_mean(tf.cast(correct, tf.float32))
  # Optimization
  opt = tf.compat.v1.train.AdamOptimizer(learning_rate=learning_rate).minimize(loss)

  # Initialize variables
  init = tf.compat.v1.global_variables_initializer()
  saver = tf.compat.v1.train.Saver()
  batch_x, batch_y = train_iterator.get_next()
  batch_test_x, batch_test_y = test_iterator.get_next()
  with tf.compat.v1.Session() as sess:
    sess.run([init])
    for i in range(train_steps):
      batch_x_value, batch_y_value = sess.run([batch_x, batch_y])
      _, loss_value = sess.run([opt, loss],
                               feed_dict={
                                   x: batch_x_value,
                                   y: batch_y_value
                               })
      if i % 100 == 0:
        tf.compat.v1.logging.info('Training step %d, loss is %f' % (i, loss_value))
      if i > 0 and i % save_every_n_steps == 0:
        accuracy_sum = 0.0
        for _ in range(eval_steps):
          test_x_value, test_y_value = sess.run([batch_test_x, batch_test_y])
          accuracy_value = sess.run(
              accuracy, feed_dict={
                  x: test_x_value,
                  y: test_y_value
              })
          accuracy_sum += accuracy_value
        tf.compat.v1.logging.info('Training step %d, accuracy is %f' %
                        (i, accuracy_sum / (eval_steps * 1.0)))
        saver.save(sess, model_dir)

def export(model, model_dir, tflite_model_file,
           use_post_training_quantize=True):
  """Export trained model to tflite model."""
  tf.compat.v1.reset_default_graph()
  x, _, output_class = model.build_model()
  saver = tf.compat.v1.train.Saver()
  sess = tf.compat.v1.Session()
  saver.restore(sess, model_dir)
  # Convert to Tflite model.
  converter = tf.compat.v1.lite.TFLiteConverter.from_session(sess, [x], [output_class])
  #converter.post_training_quantize = use_post_training_quantize
  tflite = converter.convert()
  with open(tflite_model_file, 'wb') as f:
    f.write(tflite)

def train_and_export(parsed_flags):
  """Train the MNIST LSTM model and export to TfLite."""
  model = MnistLstmModel(
      time_steps=28,
      input_size=28,
      num_lstm_layer=2,
      num_lstm_units=64,
      units=64,
      num_class=10)
  tf.compat.v1.logging.info('Starts training...')
  train(model, parsed_flags.model_dir)
  tf.compat.v1.logging.info('Finished training, starts exporting to tflite to %s ...' %
                  parsed_flags.tflite_model_file)
  export(model, parsed_flags.model_dir, parsed_flags.tflite_model_file,
         parsed_flags.use_post_training_quantize)
  tf.compat.v1.logging.info(
      'Finished exporting, model is %s' % parsed_flags.tflite_model_file)

def run_main(_):
  """Main in the TfLite LSTM tutorial."""
  parser = argparse.ArgumentParser(
      description=('Train a MNIST recognition model then export to TfLite.'))
  parser.add_argument(
      '--model_dir',
      type=str,
      help='Directory where the models will store.',
      required=True)
  parser.add_argument(
      '--tflite_model_file',
      type=str,
      help='Full filepath to the exported tflite model file.',
      required=True)
  parser.add_argument(
      '--use_post_training_quantize',
      action='store_true',
      default=True,
      help='Whether or not to use post_training_quantize.')
  parsed_flags, _ = parser.parse_known_args()
  train_and_export(parsed_flags)


def main():
  tf.compat.v1.app.run(main=run_main, argv=sys.argv[:1])


if __name__ == '__main__':
  main()

Leans heavily on backwards compatibility but can work by having the LSTM created inside a model class so that the variables are tracked. Train steps is too low for good accuracy but will run quickly for a proof of concept if you just want something to speed test on device.