Descripción general
LiteRT admite la conversión de modelos de RNN de TensorFlow en operaciones de LSTM fusionadas de LiteRT. 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 cuantificación.
Dado que existen muchas variantes de las APIs de RNN en TensorFlow, nuestro enfoque se basa en dos aspectos:
- Proporcionar compatibilidad nativa con las APIs de RNN estándar de TensorFlow, como LSTM de Keras Es la opción recomendada.
- Proporciona una interfaz en la infraestructura de conversiones para las implementaciones de RNN definidas por el usuario que se pueden conectar y convertir a LiteRT. Proporcionamos un par de ejemplos listos para usar de este tipo de conversión con las interfaces de 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 de tf-nightly pip o desde la versión principal.
Esta función de conversión está disponible cuando se convierte a LiteRT a través de un SavedModel o directamente desde el modelo de Keras. Consulta ejemplos de uso.
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()
A partir de un 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
El Colab de LSTM de Keras a LiteRT ilustra el uso integral con el intérprete de LiteRT.
APIs de RNN de TensorFlow compatibles
Conversión de LSTM de Keras (recomendado)
Admitimos la conversión lista para usar de LSTM de Keras a LiteRT. Para obtener detalles sobre cómo funciona, consulta la interfaz LSTM de Keras y la lógica de conversión aquí.
También es importante destacar el contrato de la LSTM de LiteRT con respecto a la definición de la operación de Keras:
- La dimensión 0 del tensor de entrada es el tamaño del lote.
- La dimensión 0 del tensor recurrent_weight es la cantidad de salidas.
- Se transponen los tensores weight y recurrent_kernel.
- Los tensores de peso transpuesto, recurrent_kernel transpuesto y bias se dividen en 4 tensores de igual tamaño a lo largo de la dimensión 0. Estos corresponden a la puerta de entrada, la puerta de olvido, la celda y la puerta de salida.
Variantes de LSTM de Keras
Hora principal
Los usuarios pueden elegir si quieren que la dimensión temporal sea la principal o no. La LSTM de Keras agrega un atributo time-major en los atributos de la definición de la función. Para la LSTM de secuencia unidireccional, podemos simplemente asignar el atributo principal de tiempo de unidireccional_sequence_lstm.
LSTM bidireccional
La LSTM bidireccional se puede implementar con dos capas LSTM de Keras, una para la dirección hacia adelante y otra para la dirección hacia atrás. Consulta los ejemplos aquí. Una vez que vemos el atributo go_backward, lo reconocemos como LSTM hacia atrás y, luego, agrupamos la LSTM hacia adelante y hacia atrás. Este es un trabajo futuro. Actualmente, esto crea dos operaciones UnidirectionalSequenceLSTM en el modelo de LiteRT.
Ejemplos de conversiones de LSTM definidas por el usuario
LiteRT también proporciona una forma de convertir las implementaciones de LSTM definidas por el usuario. Aquí usamos la LSTM de Lingvo como ejemplo de cómo se puede implementar. Para obtener más detalles, consulta la interfaz lingvo.LSTMCellSimple y la lógica de conversión aquí. También proporcionamos un ejemplo para otra de las definiciones de LSTM de Lingvo en la interfaz lingvo.LayerNormalizedLSTMCellSimple y su lógica de conversión aquí.
"Bring your own TensorFlow RNN" a LiteRT
Si la interfaz de la RNN de un usuario es diferente de las interfaces estándar admitidas, existen algunas opciones:
Opción 1: Escribe código de adaptador en Python de TensorFlow para adaptar la interfaz de RNN a la interfaz de RNN de Keras. Esto significa que una tf.function con anotación tf_implements en la función de la interfaz de RNN generada es idéntica a la que genera la capa LSTM de Keras. Después de esto, funcionará la misma API de conversión que se usó para la LSTM de Keras.
Opción 2: Si lo anterior no es posible (p.ej., la LSTM de Keras no tiene alguna funcionalidad que actualmente expone la op de LSTM fusionada de LiteRT, como la normalización de capas), extiende el convertidor de LiteRT escribiendo código de conversión personalizado y conéctalo al pase de MLIR prepare-composite-functions aquí. La interfaz de la función debe tratarse como un contrato de API y debe contener los argumentos necesarios para convertir las operaciones de LSTM de LiteRT fusionadas, es decir, entrada, sesgo, pesos, proyección, normalización de capas, etcétera. 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 de código de conversión que puede suponer que estos tensores son RankedTensorType y ayuda a transformarlos en tensores clasificados que corresponden a los operandos del operador de LiteRT fusionado.
Un ejemplo completo de este flujo de conversión es la conversión de LSTMCellSimple de Lingvo a LiteRT.
La LSTMCellSimple en Lingvo se define aquí. Los modelos entrenados con esta celda LSTM se pueden convertir a LiteRT de la siguiente manera:
- Encapsula todos los usos de LSTMCellSimple en un 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.
Extiende el paso prepare-composite-functions para conectar una operación compuesta personalizada a la conversión de la operación de LSTM fusionada de LiteRT. Consulta el código de conversión de LSTMCellSimple.
El contrato de conversión:
Los tensores de peso y proyección se transponen.
Los {input, recurrent} para {cell, input gate, forget gate, output gate} se extraen segmentando el tensor de peso transpuesto.
El sesgo {bias} hacia {cell, input gate, forget gate, output gate} se extrae segmentando el tensor de sesgo.
La proyección se extrae segmentando el tensor de proyección transpuesto.
Se escribe una conversión similar para LayerNormalizedLSTMCellSimple.
Se puede reutilizar el resto de la infraestructura de conversión de LiteRT, incluidos todos los pases de MLIR definidos, así como la exportación final a FlatBuffer de LiteRT.
Problemas y limitaciones conocidos
- Actualmente, solo se admite la conversión de LSTM de Keras sin estado (comportamiento predeterminado en Keras). La conversión de LSTM de Keras con estado es un trabajo futuro.
- Aún es posible modelar una capa LSTM de Keras con estado usando la capa LSTM de Keras subyacente sin estado y administrando el estado de forma explícita en el programa del usuario. Sin embargo, este programa de TensorFlow se puede convertir a LiteRT con la función que se describe aquí.
- Actualmente, el LSTM bidireccional se modela como dos operaciones UnidirectionalSequenceLSTM en LiteRT. Se reemplazará por una sola operación BidirectionalSequenceLSTM.