Conversione di TensorFlow RNN in TensorFlow Lite

Panoramica

TensorFlow Lite supporta la conversione dei modelli RNN di TensorFlow nelle operazioni LSTM fuse di TensorFlow Lite. Le operazioni fuse esistono per massimizzare le prestazioni delle implementazioni del kernel sottostanti, oltre a fornire un'interfaccia di livello superiore per definire trasformazioni complesse come la quantizzazione.

Poiché in TensorFlow sono presenti molte varianti delle API RNN, il nostro approccio è stato duplice:

  1. Fornisci supporto nativo per le API TensorFlow RNN standard come Keras LSTM. Questa è l'opzione consigliata.
  2. Fornisci un'interfaccia nell'infrastruttura di conversione per le implementazioni RNN definite dall'utente da collegare e convertire in TensorFlow Lite. Forniamo un paio di esempi pronti all'uso di tale conversione utilizzando le interfacce RNN LSTMCellSimple e LayerNormalizedLSTMCellSimple di lingvo.

API del convertitore

La funzionalità fa parte della versione TensorFlow 2.3. È disponibile anche tramite tf-nightly pip o da head.

Questa funzionalità di conversione è disponibile durante la conversione a TensorFlow Lite tramite SavedModel o direttamente dal modello Keras. Vedi esempi di utilizzo.

Dal modello salvato

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

Dal modello Keras

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

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

Esempio

Keras LSTM to TensorFlow Lite Colab illustra l'utilizzo end-to-end con l'interprete TensorFlow Lite.

API TensorFlow RNN supportate

Supportiamo la conversione immediata di Keras LSTM in TensorFlow Lite. Per i dettagli su come funziona, fare riferimento all'interfaccia Keras LSTM e alla logica di conversione qui .

Importante è anche evidenziare il contratto LSTM di TensorFlow Lite rispetto alla definizione dell'operazione Keras:

  1. La dimensione 0 del tensore di input è la dimensione del batch.
  2. La dimensione 0 del tensore recurrent_weight è il numero di uscite.
  3. I tensori peso e recurrent_kernel vengono trasposti.
  4. I tensori peso trasposto, recurrent_kernel trasposto e bias sono suddivisi in 4 tensori di uguali dimensioni lungo la dimensione 0. Questi corrispondono a porta di input, porta di dimenticanza, cella e porta di uscita .

Varianti Keras LSTM

Tempo maggiore

Gli utenti possono scegliere il tempo maggiore o nessun tempo maggiore. Keras LSTM aggiunge un attributo time-major negli attributi della funzione def. Per la sequenza unidirezionale LSTM, possiamo semplicemente mappare l'attributo time major di unidirecional_sequence_lstm.

LSTM bidirezionale

L'LSTM bidirezionale può essere implementato con due livelli Keras LSTM, uno per l'avanzamento e uno per l'indietro, vedere gli esempi qui . Una volta che vediamo l'attributo go_backward, lo riconosciamo come LSTM all'indietro, quindi raggruppiamo insieme LSTM in avanti e all'indietro. Questo è il lavoro futuro. Attualmente, questo crea due operazioni UnidirezionaliSequenceLSTM nel modello TensorFlow Lite.

Esempi di conversione LSTM definiti dall'utente

TensorFlow Lite fornisce anche un modo per convertire implementazioni LSTM definite dall'utente. Qui utilizziamo LSTM di Lingvo come esempio di come ciò possa essere implementato. Per i dettagli fare riferimento all'interfaccia lingvo.LSTMCellSimple e alla logica di conversione qui . Forniamo anche un esempio per un'altra delle definizioni LSTM di Lingvo nell'interfaccia lingvo.LayerNormalizedLSTMCellSimple e la sua logica di conversione qui .

"Porta il tuo TensorFlow RNN" su TensorFlow Lite

Se l'interfaccia RNN di un utente è diversa da quelle supportate standard, sono disponibili un paio di opzioni:

Opzione 1: scrivere il codice dell'adattatore in TensorFlow Python per adattare l'interfaccia RNN all'interfaccia Keras RNN. Ciò significa una tf.function con l'annotazione tf_implements sulla funzione dell'interfaccia RNN generata che è identica a quella generata dal livello Keras LSTM. Successivamente, funzionerà la stessa API di conversione utilizzata per Keras LSTM.

Opzione 2: Se quanto sopra non è possibile (ad esempio, in Keras LSTM mancano alcune funzionalità attualmente esposte dall'operazione LSTM fusa di TensorFlow Lite come la normalizzazione del livello), quindi estendere il convertitore TensorFlow Lite scrivendo il codice di conversione personalizzato e collegarlo alla preparazione -composite-functions MLIR-passa qui . L'interfaccia della funzione dovrebbe essere trattata come un contratto API e dovrebbe contenere gli argomenti necessari per la conversione in operazioni TensorFlow Lite LSTM fuse, ovvero input, bias, pesi, proiezione, normalizzazione dei livelli, ecc. È preferibile che i tensori vengano passati come argomenti a questo funzione per avere un rango noto (cioè RankTensorType in MLIR). Ciò rende molto più semplice scrivere codice di conversione che può assumere questi tensori come RankTensorType e aiuta a trasformarli in tensori classificati corrispondenti agli operandi dell'operatore TensorFlow Lite fuso.

Un esempio completo di tale flusso di conversione è la conversione da LSTMCellSimple a TensorFlow Lite di Lingvo.

Il LSTMCellSimple in Lingvo è definito qui . I modelli addestrati con questa cella LSTM possono essere convertiti in TensorFlow Lite come segue:

  1. Racchiudi tutti gli usi di LSTMCellSimple in una tf.function con un'annotazione tf_implements etichettata come tale (ad esempio lingvo.LSTMCellSimple sarebbe un buon nome di annotazione qui). Assicurati che la tf.function generata corrisponda all'interfaccia della funzione prevista nel codice di conversione. Questo è un contratto tra l'autore del modello che aggiunge l'annotazione e il codice di conversione.
  2. Estendi il passaggio delle funzioni prepare-composite per collegare un'operazione composita personalizzata alla conversione operativa LSTM fusa di TensorFlow Lite. Vedere il codice di conversione LSTMCellSimple .

    Il contratto di conversione:

  3. I tensori peso e proiezione sono trasposti.

  4. L '{input, ricorrente} in {cell, input gate, dimenticato gate, output gate} viene estratto affettando il tensore del peso trasposto.

  5. Il {bias} su {cella, porta di ingresso, porta di dimenticanza, porta di uscita} viene estratto affettando il tensore di polarizzazione.

  6. La proiezione viene estratta affettando il tensore di proiezione trasposto.

  7. Una conversione simile viene scritta per LayerNormalizedLSTMCellSimple .

  8. Il resto dell'infrastruttura di conversione di TensorFlow Lite, inclusi tutti i passaggi MLIR definiti e l'esportazione finale nel flatbuffer di TensorFlow Lite, può essere riutilizzato.

Problemi/limitazioni noti

  1. Attualmente è disponibile il supporto solo per la conversione di Keras LSTM senza stato (comportamento predefinito in Keras). La conversione Stateful Keras LSTM è un lavoro futuro.
  2. È ancora possibile modellare un livello Keras LSTM con stato utilizzando il livello Keras LSTM senza stato sottostante e gestendo lo stato esplicitamente nel programma utente. Un programma TensorFlow di questo tipo può ancora essere convertito in TensorFlow Lite utilizzando la funzionalità descritta qui.
  3. L'LSTM bidirezionale è attualmente modellato come due operazioni UnidirectionSequenceLSTM in TensorFlow Lite. Questo sarà sostituito con un singolo BidirezionaleSequenceLSTM op.