Conversión de TensorFlow RNN a TensorFlow Lite

Visión de conjunto

TensorFlow Lite admite la conversión de modelos RNN de TensorFlow a operaciones LSTM fusionadas 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 cuantificación.

Dado que hay muchas variantes de las API de RNN en TensorFlow, nuestro enfoque ha sido doble:

  1. Proporcione soporte nativo para las API RNN de TensorFlow estándar como Keras LSTM. Esta es la opción recomendada.
  2. Proporcione una interfaz en la infraestructura de conversión para que las implementaciones de RNN definidas por el usuario se conecten y se conviertan a TensorFlow Lite. Proporcionamos un par de ejemplos listos para usar de dicha conversión utilizando las interfaces LSTMCellSimple y LayerNormalizedLSTMCellSimple RNN de lingvo.

API de convertidor

La característica es parte de la versión TensorFlow 2.3. También está disponible a través del tf-nightly pip o desde head.

Esta funcionalidad 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. Ver usos de ejemplo.

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

Se admiten las API RNN de TensorFlow

Admitimos la conversión lista para usar de Keras LSTM a TensorFlow Lite. Para obtener detalles sobre cómo funciona, consulte la interfaz Keras LSTM ya la lógica de conversión aquí .

También es importante destacar el contrato LSTM de TensorFlow Lite con respecto a la definición de 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 de peso_recurrente es el número de salidas.
  3. Los tensores de peso y núcleo recurrente se transponen.
  4. Los tensores de peso transpuesto, núcleo recurrente transpuesto y polarización 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 Keras LSTM

tiempo mayor

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

LSTM bidireccional

El LSTM bidireccional se puede implementar con dos capas Keras LSTM, una para avanzar y otra para retroceder; consulte los ejemplos aquí . Una vez que vemos el atributo go_backward, lo reconocemos como LSTM hacia atrás, luego agrupamos los LSTM hacia adelante y hacia atrás. Este es un trabajo futuro. Actualmente, esto crea dos operaciones UnidireccionalSequenceLSTM en el modelo TensorFlow Lite.

Ejemplos de conversión de LSTM definidos por el usuario

TensorFlow Lite también proporciona una forma de convertir implementaciones de LSTM definidas por el usuario. Aquí usamos el LSTM de Lingvo como ejemplo de cómo se puede implementar. Para obtener más información, consulte 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í .

"Traiga su propio TensorFlow RNN" a TensorFlow Lite

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

Opción 1: escriba el código del adaptador en TensorFlow python para adaptar la interfaz RNN a la interfaz Keras RNN. Esto significa una tf.function con anotación tf_implements en la función de la interfaz RNN generada que es idéntica a la generada por la capa Keras LSTM. Después de esto, funcionará la misma API de conversión utilizada para Keras LSTM.

Opción 2: si lo anterior no es posible (p. ej., Keras LSTM no tiene alguna funcionalidad que actualmente está expuesta por la normalización de capa similar a la operación LSTM fusionada de TensorFlow Lite), extienda el convertidor de TensorFlow Lite escribiendo un código de conversión personalizado y conéctelo a la preparación -funciones compuestas MLIR-pase aquí . La interfaz de la función debe tratarse como un contrato de API y debe contener los argumentos necesarios para convertir a operaciones fusionadas de TensorFlow Lite LSTM, es decir, entrada, sesgo, pesos, proyección, normalización de capas, etc. Es preferible que los tensores se pasen como argumentos a este función para tener rango conocido (es decir, RankedTensorType en MLIR). Esto hace que sea mucho más fácil escribir código de conversión que pueda asumir estos tensores como RankedTensorType y ayuda a transformarlos en tensores clasificados correspondientes a los operandos fusionados del operador TensorFlow Lite.

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

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. Envuelva todos los usos de LSTMCellSimple en una tf.function con una anotación tf_implements etiquetada como tal (por ejemplo, lingvo.LSTMCellSimple sería un buen nombre de anotación aquí). Asegúrese de que la función tf. 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. Extienda el pase de preparación de funciones compuestas para conectar una operación compuesta personalizada a la conversión de operación LSTM fusionada de TensorFlow Lite. Consulte el código de conversión LSTMCellSimple .

    El contrato de conversión:

  3. Los tensores de peso y proyección se transponen.

  4. La {entrada, recurrente} a {celda, puerta de entrada, puerta de olvido, puerta de salida} se extraen cortando el tensor de peso transpuesto.

  5. El {sesgo} a {celda, puerta de entrada, puerta de olvido, puerta de salida} se extraen cortando el tensor de polarización.

  6. La proyección se extrae cortando el tensor de proyección transpuesto.

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

  8. El resto de la infraestructura de conversión de TensorFlow Lite, incluidos todos los pases de MLIR definidos, así como la exportación final al búfer plano de TensorFlow Lite, se pueden reutilizar.

Problemas conocidos/limitaciones

  1. Actualmente solo hay soporte para convertir Keras LSTM sin estado (comportamiento predeterminado en Keras). La conversión Stateful Keras LSTM es un trabajo futuro.
  2. Todavía es posible modelar una capa LSTM de Keras con estado utilizando la capa LSTM de Keras sin estado subyacente y administrando el estado explícitamente en el programa de usuario. Tal programa TensorFlow todavía se puede convertir a TensorFlow Lite usando la función que se describe aquí.
  3. El LSTM bidireccional se modela actualmente como dos operaciones UnidireccionalSequenceLSTM en TensorFlow Lite. Esto se reemplazará con una sola operación BidireccionalSequenceLSTM.