Conversión de RNN de TensorFlow a TensorFlow Lite

Descripción general

TensorFlow Lite es compatible con la conversión de modelos de RNN de TensorFlow en operaciones fusionadas con LSTM de TensorFlow Lite. Las operaciones fusionadas existen para maximizar el rendimiento de sus implementaciones de kernel subyacentes, así como para proporcionar una interfaz de nivel superior para definir transformaciones complejas como la cuantización.

Dado que hay muchas variantes de las APIs de RNN en TensorFlow, nuestro enfoque se ha dobletado:

  1. Proporciona compatibilidad nativa con las APIs de RNN estándares de TensorFlow, como Keras LSTM. Es la opción recomendada.
  2. Proporciona una interfaz en la infraestructura de conversiones de las implementaciones de RNN definidas por el usuario para conectarse y convertirse a TensorFlow Lite. Proporcionamos un par de ejemplos listos para usar de esa conversión con las interfaces RNN LSTMCellSimple y LayerNormalizedLSTMCellSimple de Lingvo.

API de Converter

Esta función forma parte de la versión 2.3 de TensorFlow. También está disponible a través del pip tf-nightly o desde head.

Esta función de conversión está disponible cuando se convierte a TensorFlow Lite a través de un modelo guardado o directamente desde el modelo de Keras. Consulta ejemplos de uso.

A partir del modelo guardado

# build a saved model. Here concrete_function is the exported function
# corresponding to the TensorFlow model containing one or more
# Keras LSTM layers.
saved_model, saved_model_dir = build_saved_model_lstm(...)
saved_model.save(saved_model_dir, save_format="tf", signatures=concrete_func)

# Convert the model.
converter = TFLiteConverter.from_saved_model(saved_model_dir)
tflite_model = converter.convert()

Del modelo de Keras

# build a Keras model
keras_model = build_keras_lstm(...)

# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()

Ejemplo

Keras LSTM a TensorFlow Lite Colab ilustra el uso de extremo a extremo con el intérprete de TensorFlow Lite.

Compatible con las APIs de RNN de TensorFlow

Admitimos la conversión lista para usar de Keras LSTM a TensorFlow Lite. Si deseas obtener detalles sobre cómo funciona, consulta la interfaz de LSTM de Keras y la lógica de conversión aquí.

También es importante destacar el contrato de LSTM de TensorFlow Lite con respecto a la definición de la operación de Keras:

  1. La dimensión 0 del tensor input es el tamaño del lote.
  2. La dimensión 0 del tensor recurrent_weight es la cantidad de salidas.
  3. Se transponen los tensores weight y recurrent_kernel.
  4. Los tensores de peso transpuesto, recurrentet_kernel transpuesto y bias se dividen en 4 tensores de igual tamaño a lo largo de la dimensión 0. Estas corresponden a puerta de entrada, puerta de olvido, celda y puerta de salida.

Variantes de Keras LSTM

Mejor horario

Los usuarios pueden elegir una mayor cantidad de tiempo o ninguna. Keras LSTM agrega un atributo mayor de tiempo en los atributos de definición de la función. Para la secuencia unidireccional de LSTM, simplemente podemos asignar al atributo tiempo principal de unidirecional_Sequence_lstm.

LSTM bidireccional

La LSTM bidireccional se puede implementar con dos capas de LSTM de Keras, una para avanzar y otra para atrás. Consulta ejemplos aquí. Una vez que vemos el atributo go_backward, lo reconocemos como LSTM hacia atrás y, luego, agrupamos la LSTM hacia delante y hacia atrás. Este es un trabajo futuro. Actualmente, esto crea dos operaciones UnidirectionalSequenceLSTM en el modelo de TensorFlow Lite.

Ejemplos de conversiones de LSTM definidas por el usuario

TensorFlow Lite también proporciona una forma de convertir implementaciones de LSTM definidas por los usuarios. Aquí, usamos la LSTM de Lingvo como ejemplo de cómo se puede implementar. Para obtener más información, consulta la interfaz de lingvo.LSTMCellSimple y la lógica de conversión aquí. También proporcionamos un ejemplo de otra de las definiciones de LSTM de Lingvo en la interfaz de lingvo.LayerNormalizedLSTMCellSimple y su lógica de conversión aquí.

“Trae tu propia RNN de TensorFlow” a TensorFlow Lite

Si la interfaz RNN de un usuario es diferente de las estándar compatibles, existen algunas opciones:

Opción 1: Escribe el código del adaptador en TensorFlow para Python para adaptar la interfaz de RNN a la interfaz de RNN de Keras. Esto significa una tf.function con la anotación tf_implements en la función de la interfaz RNN generada, que es idéntica a la que genera la capa de LSTM de Keras. Luego, funcionará la misma API de conversión que se usa para Keras LSTM.

Opción 2: Si no es posible lo anterior (p.ej., a la LSTM de Keras le falta alguna funcionalidad que actualmente está expuesta por la op de LSTM fusionada de TensorFlow Lite, como la normalización de capas), luego extiende el conversor de TensorFlow Lite escribiendo un código de conversión personalizado y conéctalo a las funciones de preparación y composición de MLIR aquí. La interfaz de la función debe tratarse como un contrato de API y debe contener los argumentos necesarios para convertir en operaciones fusionadas de LSTM de TensorFlow Lite, es decir, entrada, sesgo, pesos, proyección, normalización de capas, etc. Es preferible que los tensores que se pasan como argumentos a esta función tengan un rango conocido (es decir, RankedTensorType en MLIR). Esto facilita mucho la escritura del código de conversión que puede asumir estos tensores como RankedTensorType y ayuda a transformarlos en tensores clasificados que corresponden a los operandos del operador fusionado de TensorFlow Lite.

Un ejemplo completo de ese flujo de conversión es la conversión de LSTMCellSimple de Lingvo a TensorFlow Lite.

El LSTMCellSimple en Lingvo se define aquí. Los modelos entrenados con esta celda LSTM se pueden convertir a TensorFlow Lite de la siguiente manera:

  1. Une todos los usos de LSTMCellSimple en una tf.function con una anotación tf_implements que esté etiquetada como tal (p.ej., lingvo.LSTMCellSimple sería un buen nombre de anotación aquí). Asegúrate de que la tf.function que se genera coincida con la interfaz de la función esperada en el código de conversión. Este es un contrato entre el autor del modelo que agrega la anotación y el código de conversión.
  2. Extiende el pase de funciones compuestas de preparación para conectar una op compuesta personalizada a la conversión de op de LSTM fusionada de TensorFlow Lite. Consulta el código de conversión LSTMCellSimple.

    El contrato de conversión tiene las siguientes características:

  3. Se transponen los tensores Weight y projection.

  4. Los valores de {input, recurrentet} a {cell, input gate, ignore gate, output gate} se extraen mediante la división del tensor de peso transpuesto.

  5. Los valores de {bias} a {cell, input gate, ignore gate, output gate} se extraen mediante la división del tensor de sesgo.

  6. La proyección se extrae mediante la división del tensor de proyección transpuesto.

  7. Se escribió una conversión similar para LayerNormalizedLSTMCellSimple.

  8. Se pueden volver a usar el resto de la infraestructura de conversiones de TensorFlow Lite, incluidos todos los pases de MLIR definidos y la exportación final al búfer plano de TensorFlow Lite.

Limitaciones o problemas conocidos

  1. Actualmente, solo se admite la conversión de Keras LSTM sin estado (comportamiento predeterminado en Keras). La conversión de Keras con estado en LSTM es un trabajo futuro.
  2. Aún es posible modelar una capa de Keras LSTM con estado mediante la capa de Keras LSTM sin estado subyacente y administrar el estado de forma explícita en el programa de usuario. Este programa de TensorFlow aún se puede convertir a TensorFlow Lite con la función que se describe aquí.
  3. Actualmente, la LSTM bidireccional se modela como dos operaciones UnidirectionalSequenceLSTM en TensorFlow Lite. Esto se reemplazará por una sola op BidirectionalSequenceLSTM.