5.2 KiB
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:
- tf.nn.rnn_cell.LSTMCell
- tf.nn.rnn_cell.RNNCell
- tf.nn.rnn_cell.GRUCell
- tf.nn.rnn_cell.BasicLSTMCell
- tf.nn.rnn_cell.BasicRNNCell
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:
- tf.nn.static_rnn
where a
sequence_lengthis specified - tf.nn.dynamic_rnn
- tf.nn.bidirectional_dynamic_rnn
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:
- tf.lite.experimental.nn.dynamic_rnn
- replacement for tf.nn.dynamic_rnn
- tf.lite.experimental.nn.bidirectional_dynamic_rnn
- replacement for tf.nn.bidirectional_dynamic_rnn
- tf.lite.experimental.nn.TfLiteRNNCell
- replacement for tf.nn.rnn_cell.RNNCell
- tf.lite.experimental.nn.TfLiteLSTMCell
- replacement for tf.nn.rnn_cell.LSTMCell
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.