Conversión de RNN de TensorFlow a LiteRT

Descripción general

LiteRT admite la conversión de los modelos de RNN de TensorFlow a los de LiteRT. operaciones de LSTM combinadas. Las operaciones fusionadas existen para maximizar el rendimiento de sus implementaciones de kernel subyacentes y proporcionan un nivel para definir transformaciones complejas, como la cuantización.

Debido a que existen muchas variantes de APIs de RNN en TensorFlow, nuestro enfoque ha sido doble:

  1. Proporciona compatibilidad nativa para las APIs de RNN de TensorFlow estándar, como Keras LSTM. Es la opción recomendada.
  2. Proporcionar una interfaz en la infraestructura de conversiones Implementaciones de RNN definidas por el usuario para conectarse y convertirse en LiteRT. Ofrecemos un par de ejemplos listos para usar de tales de conversiones con lenguaje LSTMCellSimple y LayerNormalizedLSTMCellSimple RNN.

API de Converter

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

Esta función de conversión estará disponible cuando realices la conversión a LiteRT. a través de un modelo guardado o desde el modelo de Keras directamente. Consulta ejemplos de usos.

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()

Desde el 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 LiteRT Colab ilustra el uso de extremo a extremo con el intérprete de LiteRT.

APIs de RNN de TensorFlow compatibles

Admitimos la conversión lista para usar de Keras LSTM a LiteRT. Para sobre el funcionamiento de esta función, consulte la Interfaz de Keras LSTM y en la lógica de conversión. aquí.

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

  1. La dimensión 0 del tensor de entrada es el tamaño del lote.
  2. La dimensión 0 del tensor recurrent_weight es el número de de salida.
  3. Los tensores weight y recurrent_kernel se transponen.
  4. Los tensores de peso transpuesto, recurrent_kernel transpuesto y sesgo se dividida en 4 tensores de igual tamaño a lo largo de la dimensión 0. Estas corresponden puerta de entrada, puerta de descarte, celda y puerta de salida.

Variantes de Keras LSTM

Principal

Los usuarios pueden elegir tiempo principal o sin tiempo mayor. Keras LSTM agrega un tiempo mayor en los atributos def de la función. Para la secuencia unidireccional de LSTM, puede asignarse a unidirecional_Sequence_lstm atributo mayor de tiempo.

LSTM bidireccional

El LSTM bidireccional se puede implementar con dos capas de Keras LSTM, una para hacia adelante y una hacia atrás, ve ejemplos aquí. Una vez que vemos el atributo go_backward, lo reconocemos como LSTM hacia atrás y, luego, agrupamos en orden y y revertir LSTM juntos. Este es un trabajo futuro. Actualmente, Esto crea dos operaciones UnidirectionalSequenceLSTM en LiteRT. un modelo de responsabilidad compartida.

Ejemplos de conversión de LSTM definidas por el usuario

LiteRT también proporciona una forma de convertir las LSTM definidas por el usuario de Google Cloud. Aquí usamos la LSTM de Lingvo como ejemplo de cómo se puede cuando se implementa un plan. Para obtener más información, 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 Interfaz lingvo.LayerNormalizedLSTMCellSimple y su lógica de conversiones aquí.

“Trae tu propia RNN de TensorFlow” a LiteRT

Si la interfaz RNN de un usuario es diferente de las admitidas estándar, hay hay un par de opciones:

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

Opción 2: Si no es posible realizar lo anterior (p.ej., a Keras LSTM le faltan algunos que está expuesta actualmente por la op fusionada de LSTM de LiteRT, como la normalización de capas) y, luego, extiende el conversor de LiteRT escribiendo código de conversión personalizado y conectarlo a las funciones prepare-composite-functions MLIR-pass aquí. La interfaz de la función debe tratarse como un contrato de API y debe contiene los argumentos necesarios para convertir a LSTM de LiteRT fusionado operaciones, es decir, entrada, sesgo, ponderación, proyección, normalización de capas, etc. es preferible que los tensores pasados como argumentos a esta función conozcan (p.ej., RankedTensorType en MLIR). Esto hace que sea mucho más fácil escribir de conversión que puede asumir estos tensores como RankedTensorType y ayuda transformarlos en tensores clasificados correspondientes al modelo LiteRT fusionado operandos del operador.

Un ejemplo completo de este flujo de conversión es LSTMCellSimple de Lingvo para Conversión de LiteRT.

Se define la LSTMCellSimple en Lingvo aquí. Los modelos entrenados con esta celda de LSTM se pueden convertir a LiteRT de la siguiente manera: sigue:

  1. Unir todos los usos de LSTMCellSimple en una tf.function con un tf_implements anotación etiquetada como tal (por ejemplo, lingvo.LSTMCellSimple sería una nombre de anotación). Asegúrate de que la función tf.function que se genera Coincide con la interfaz de la función que se espera en el código de conversión. Esta es un contrato entre el autor del modelo que agrega la anotación y el el código de conversión.
  2. Extiende el pase de preparación de funciones compuestas para conectar una op personalizada compuesta en la conversión de operaciones de LSTM combinada con LiteRT. Consulta LSTMCellSimple el código de conversión.

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

  3. Los tensores de Weight y projection se transponen.

  4. La {input, recurrente} como {cell, input gate, exit gate, output gate} se cortan el tensor de peso transpuesto.

  5. Los valores de {bias} a {bias} son extrayendo el tensor de sesgo.

  6. Se extrae la proyección fragmentando el tensor de proyección transpuesto.

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

  8. El resto de la infraestructura de conversiones de LiteRT, incluidos Pases del MLIR definido, así como la exportación final al búfer plano de LiteRT reutilizados.

Limitaciones o problemas conocidos

  1. Actualmente, solo se admite la conversión de LSTM de Keras sin estado (configuración predeterminada en Keras). La conversión de Keras LSTM con estado es un trabajo futuro.
  2. Aún es posible modelar una capa de Keras LSTM con estado usando el la capa subyacente de LSTM de Keras sin estado y la administración del estado explícitamente en el programa de usuarios. Este programa de TensorFlow aún puede convertirse en LiteRT que usa la función que se describe aquí.
  3. Actualmente, la LSTM bidireccional está modelada como dos UnidirectionalSequenceLSTM. operaciones en LiteRT. Esta se reemplazará con una sola Operación BidirectionalSequenceLSTM.