Convert RNN models

The TensorFlow Lite interpreter currently implements a subset of TensorFlow operations, meaning some model architectures cannot immediately be converted due to missing operations.

Some RNN-based architectures are affected by this. The following document outlines the current state of play and provides strategies for converting RNN models.

Currently supported

Currently, RNN models using tf.nn.static_rnn can be converted successfully as long as no sequence_length is specified.

The following tf.nn.rnn_cell operations work with tf.nn.static_rnn:

In addition, TensorFlow Lite provides some experimental drop-in replacements for RNN operations that enable dynamic RNN architectures with TensorFlow Lite.

Drop-in replacements are available for the following:

Not currently supported

TensorFlow Lite does not currently support Control Flow operations. This means that, unless one of the conversion strategies discussed in the next section are employed, models built with the following TensorFlow functions will not convert successfully:

Note: TensorFlow Lite plans to implement all required Control Flow operations by the end of 2019. At this point, all RNN architectures will convert successfully.

Conversion strategies

To convert an RNN model that uses the functions specified above, you will have to modify its architecture and retrain it. The following strategies can be used.

1. Refactoring

The simplest approach, if possible, is to refactor the model architecture to use tf.nn.static_rnn without sequence_length.

2. Drop-in replacements that use op hints and fused ops

TensorFlow Lite provides the some experimental drop-in replacements for RNN operations that enable dynamic RNN architectures with TensorFlow Lite. Using OpHints, they run normally during training, but are substituted with special fused ops when run by the Lite interpreter.

The following drop-in replacements are available:

Note: These replacements must be used together. For example, if you are using tf.lite.experimental.nn.dynamic_rnn, you must combine it with tf.lite.experimental.nn.TfLiteRNNCell instead of using tf.nn.rnn_cell.RNNCell.

Instead of tf.nn.rnn_cell.MultiRNNCell, you should use tf.keras.layers.StackedRNNCells.

For a tutorial on using these replacements, see TensorFlow Lite LSTM ops API.

For a Colab demonstrating these classes, refer to TensorFlowLite_LSTM_Keras_Tutorial.

Note: There is no replacement available for tf.nn.rnn_cell.GRUCell.