Conversión de TensorFlow RNN a TensorFlow Lite

Descripción general

TensorFlow Lite admite la conversión de modelos RNN de TensorFlow a operaciones LSTM fusionadas de TensorFlow Lite. Existen operaciones fusionadas 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 existen muchas variantes de API RNN en TensorFlow, nuestro enfoque ha sido doble:

  1. Proporcione soporte nativo para las API RNN estándar de TensorFlow como Keras LSTM. Esta es la opción recomendada.
  2. Proporcione una interfaz en la infraestructura de conversión para que las implementaciones 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 del convertidor

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

Esta funcionalidad de conversión está disponible al convertir a TensorFlow Lite a través de un modelo guardado o directamente desde el modelo de Keras. Vea 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 un extremo a otro con el intérprete de TensorFlow Lite.

API de RNN de TensorFlow compatibles

Admitimos la conversión inmediata de Keras LSTM a TensorFlow Lite. Para obtener detalles sobre cómo funciona esto, consulte la interfaz Keras LSTM. y a 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 recurrent_weight es el número de salidas.
  3. Los tensores de peso y kernel_recurrente se transponen.
  4. Los tensores de peso transpuesto, núcleo recurrente transpuesto y sesgo 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 principal o no tiempo principal. Keras LSTM agrega un atributo de tiempo principal en los atributos def de la función. Para la secuencia unidireccional LSTM, simplemente podemos asignar el atributo principal de tiempo de unidirecional_sequence_lstm.

LSTM bidireccional

El LSTM bidireccional se puede implementar con dos capas de Keras LSTM, una para adelante y otra para atrás; consulte los ejemplos aquí . Una vez que vemos el atributo go_backward, lo reconocemos como LSTM hacia atrás, luego agrupamos 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 LSTM definidos por el usuario

TensorFlow Lite también proporciona una forma de convertir implementaciones de LSTM definidas por el usuario. Aquí utilizamos 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 de otra de las definiciones LSTM de Lingvo en la interfaz lingvo.LayerNormalizedLSTMCellSimple y su lógica de conversión aquí .

"Traiga su propio RNN de TensorFlow" a TensorFlow Lite

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

Opción 1: escribir código de adaptador en TensorFlow Python para adaptar la interfaz RNN a la interfaz Keras RNN. Esto significa una función tf.function con la 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 (por ejemplo, a Keras LSTM le falta 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 al proceso de preparación. -funciones-compuestas MLIR-pase aquí . La interfaz de la función debe tratarse como un contrato API y debe contener los argumentos necesarios para convertir a operaciones LSTM fusionadas de TensorFlow Lite, es decir, entrada, polarización, pesos, proyección, normalización de capas, etc. Es preferible que los tensores pasados ​​como argumentos a esto función para tener un rango conocido (es decir, ClasificadoTensorType en MLIR). Esto hace que sea mucho más fácil escribir código de conversión que pueda asumir estos tensores como ClasificadoTensorType y ayuda a transformarlos en tensores clasificados correspondientes a los operandos del operador fusionado de 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 función 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. Amplíe el paso prepare-composite-functions para conectar una operación compuesta personalizada a la conversión de operaciones LSTM fusionadas de TensorFlow Lite. Consulte el código de conversión LSTMCellSimple .

    El contrato de conversión:

  3. Se transponen los tensores de peso y proyección .

  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 extrae 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 MLIR definidos, así como la exportación final al búfer plano de TensorFlow Lite, se pueden reutilizar.

Problemas/limitaciones conocidos

  1. Actualmente solo hay soporte para convertir Keras LSTM sin estado (comportamiento predeterminado en Keras). La conversión con estado de Keras LSTM es un trabajo futuro.
  2. Todavía es posible modelar una capa Keras LSTM con estado utilizando la capa Keras LSTM sin estado subyacente y administrando el estado explícitamente en el programa de usuario. Un programa de TensorFlow de este tipo aún se puede convertir a TensorFlow Lite utilizando la función que se describe aquí.
  3. Actualmente, LSTM bidireccional se modela como dos operaciones UnidireccionalSequenceLSTM en TensorFlow Lite. Esto será reemplazado por una única operación BidireccionalSequenceLSTM.