Conversione da TensorFlow RNN a TensorFlow Lite

Panoramica

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

Poiché in TensorFlow esistono 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 predefiniti di questa conversione utilizzando le interfacce RNN LSTMCellSimple e LayerNormalizedLSTMCellSimple RNN di lingvo.

API Converter

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

Questa funzionalità di conversione è disponibile quando si esegue la conversione a TensorFlow Lite tramite un modello SaveModel 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 in TensorFlow Lite Colab descrive l'utilizzo end-to-end con l'interprete di TensorFlow Lite.

API TensorFlow RNNs supportate

Supportiamo la conversione immediata di Keras LSTM a TensorFlow Lite. Per informazioni dettagliate su come funziona, consulta l'interfaccia Keras LSTM e la logica di conversione qui.

È anche importante evidenziare il contratto LSTM di TensorFlow Lite in relazione 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 output.
  3. I tensori weight e recurrent_kernel vengono trasposti.
  4. I tensori del peso trasposto, del recurrent_kernel e del bias trasposti sono suddivisi in 4 tensori di dimensioni uguali lungo la dimensione 0. Questi corrispondono a gate di input, porta di eliminazione, cella e gate di output.

Varianti Keras LSTM

Ora principale

Gli utenti possono scegliere tra "Migliore" o nessuna metrica "Principale". Keras LSTM aggiunge un attributo temporale importante negli attributi della funzione def. Per la sequenza unidirezionale LSTM, possiamo semplicemente eseguire la mappatura all'attributo temporale principale di unidirecional_ sequence_lstm.

LSTM bidirezionale

L'LSTM bidirezionale può essere implementata con due livelli Keras LSTM, uno per il passaggio avanti e uno per il passaggio a ritroso. Consulta gli esempi qui. Una volta visualizzato l'attributo go_backward, lo riconosciamo come LSTM indietro, quindi raggruppiamo LSTM avanti e indietro. Questo è un lavoro futuro. Attualmente, vengono create due operazioni UnidirectionalSequenceLSTM nel modello TensorFlow Lite.

Esempi di conversione LSTM definiti dall'utente

TensorFlow Lite fornisce anche un modo per convertire le implementazioni LSTM definite dall'utente. Qui usiamo l'LSTM di Lingvo come esempio di come può essere implementato. Per maggiori dettagli, consulta l'interfaccia lingvo.LSTMCellSimple e la logica di conversione qui. Forniamo anche un esempio per un'altra definizione LSTM di Lingvo nell'interfaccia di lingvo.LayerNormalizedLSTMCellSimple e nella sua logica di conversione qui.

"Bring your own TensorFlow RNN" su TensorFlow Lite

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

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

Opzione 2: se non è possibile farlo (ad esempio, nell'LSTM Keras manca alcune funzionalità attualmente esposte dall'operazione LSTM fusa di TensorFlow Lite per la normalizzazione dei livelli), estendi il convertitore TensorFlow Lite scrivendo un codice di conversione personalizzato e inseriscilo in MLIR-pass qui. L'interfaccia della funzione deve essere trattata come un contratto API e deve contenere gli argomenti necessari per la conversione in operazioni LSTM fuse di TensorFlow Lite, ovvero input, bias, ponderazioni, proiezione, normalizzazione del livello e così via. È preferibile che i tensori passati come argomenti a questa funzione abbiano un ranking noto (ovvero RankedTensorType in MLIR). Ciò semplifica molto la scrittura del codice di conversione che può assumere questi tensori come RankedTensorType e aiuta a trasformarli in tensori classificati corrispondenti agli operandi dell'operatore fuso di TensorFlow Lite.

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

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

  1. Aggrega tutti gli utilizzi di LSTMCellSimple in una tf.function con un'annotazione tf_implements etichettata come tale (ad esempio, lingvo.LSTMCellSimple sarebbe un nome corretto per l'annotazione qui). Assicurati che tf.function generata corrisponda all'interfaccia della funzione prevista nel codice di conversione. Si tratta di un contratto tra l'autore del modello che aggiunge l'annotazione e il codice di conversione.
  2. Estendi il passaggio delle funzioni composite di preparazione per collegare un'operazione composita personalizzata alla conversione dell'operazione LSTM con fusione di TensorFlow Lite. Consulta il codice di conversione LSTMCellSimple.

    Il contratto di conversione:

  3. I tensori di peso e di proiezione vengono trasposti.

  4. Gli elementi da {input, recurrent} a {cell, input gate, differiscono gate, output gate} vengono estratti tramite lo slicing del tensore di peso trasposto.

  5. Il valore da {bias} a {cell, input gate, hide gate, output gate} viene estratto mediante lo slicing del tensore di bias.

  6. La proiezione viene estratta dal taglio del tensore di proiezione trasposto.

  7. Una conversione simile è scritta per LayerNormalizedLSTMCellSimple.

  8. È possibile riutilizzare il resto dell'infrastruttura di conversione di TensorFlow Lite, inclusi tutti i pass MLIR definiti e l'esportazione finale nel flatbuffer TensorFlow Lite.

Problemi noti/limitazioni

  1. Attualmente è supportato solo per la conversione di Keras LSTM stateless (comportamento predefinito in Keras). La conversione Stateful Keras LSTM è un lavoro futuro.
  2. È ancora possibile modellare un livello Keras LSTM stateful utilizzando il livello sottostante Keras LSTM stateless e gestendo lo stato in modo esplicito nel programma utente. Un programma TensorFlow di questo tipo può comunque essere convertito in TensorFlow Lite utilizzando la funzionalità qui descritta.
  3. Bidirectional LSTM è attualmente modellato come due operazioni UnidirectionalSequenceLSTM in TensorFlow Lite. Questo verrà sostituito con una singola operazione BidirectionalSequenceLSTM.