RNN tflite model doesn't work when deployed on a microcontroller #1570
Description
System information
- Tensorflow version: 2.9.1
Describe the feature and the current behavior/state.
RNN tflite model not working when deployed on a microcontroller unless it is unrolled which increases the size of the model tenfold.
I want to compare different types of RNN tflite-micro models. I have created a custom RNN cell that I want to compare with the LSTM cell, GRU cell, and SimpleRNN cell.
The following code shows how the network is created and converted to tflite model
import tensorflow as tf
# create the rnn cell using one of these RNN cells
units = <value>
#RNNcell = tf.keras.layers.SimpleRNNCell(units)
#RNNcell = tf.keras.layers.GRUCell(units)
RNNcell = tf.keras.layers.LSTMCell(units)
# create the model using the cell in an RNN layer
timesteps = <value>
n_features = <value>
dropout = <value>
dense_units = <value>
outputs = <value>
model = tf.keras.Sequential([
tf.keras.layers.RNN(RNNcell, input_shape=(timesteps, n_features),
tf.keras.layers.Dense(dense_units, activation='relu'),
tf.keras.layers.Dense(outputs, activation='softmax', name="output")
])
# then create model for inference with a fixed batch size
model = tf.keras.Sequential([
tf.keras.layers.RNN(RNNcell, batch_size=1, input_shape=(timesteps, n_features),
tf.keras.layers.Dense(dense_units, activation='relu'),
tf.keras.layers.Dense(outputs, activation='softmax', name="output")
])
#train the model
EPOCHS = <value>
BATCH_SIZE = <value>
LEARNING_RATE = <value>
opt = tf.keras.optimizers.Adam(learning_rate=LEARNING_RATE)
model.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
model.fit(
train_X, train_y,
epochs=EPOCHS,
validation_data=(validation_X, validation_y),
shuffle=False,
batch_size=BATCH_SIZE)
# move weights to model with fixed batch_size
model_inference.set_weights(loaded_model.get_weights())
model_inference.compile(optimizer=opt, loss='categorical_crossentropy', metrics=['accuracy'])
# save to tflite format
converter = tf.lite.TFLiteConverter.from_keras_model(model_inference)
tfmodel = converter.convert()
open(f'model.tflite', 'wb').write(tfmodel)
## EXTRA
# test operations with tflite interpreter in python
input_data = <value>
interpreter = tf.lite.Interpreter(model_content=tfmodel)
interpreter.allocate_tensors()
_input_index = interpreter.get_input_details()[0]['index']
interpreter.set_tensor(_input_index, input_data)
interpreter.invoke()
The tflite model works in the python interpreter but does not run on a microcontroller (nRF5340) unless the RNN layer is unrolled, which increases the size of the model significantly.
This is how the model looks in Netron and the error from the interpreter on the microcontroller indicates that INT32 is not supported for ADD in the WHILE operator.
Type INT32 (2) not supported
Node ADD (number 0) failed to invoke with status 1
Node WHILE (number 1) failed to invoke with status 1
Will this change the current api? How?
don't know
Who will benefit with this feature?
Everybody who needs and thought RNNs were supported in Tensorflow Lite Micro.
Any Other info.
Similar to issue #907 but in that issue that person is unrolling.
Also, a tflite model with created with tf.keras.layers.LSTM creates the operator unidirectionalLSTM but the model is created with a tf.keras.layers.LSTMCell in a tf.keras.layers.RNN layer then it has a reshape and a while operator instead.
Is it possible to use the unidirectionalLSTM operator for the other cells as well somehow?