Panoramica
LiteRT supporta la conversione dei modelli RNN TensorFlow in operazioni LSTM fused di LiteRT. Le operazioni combinate esistono per massimizzare le prestazioni delle implementazioni del kernel sottostanti, nonché per fornire un'interfaccia di livello superiore per definire trasformazioni complesse come la quantizzazione.
Poiché in TensorFlow esistono molte varianti di API RNN, il nostro approccio è stato duplice:
- Fornire supporto nativo per le API RNN TensorFlow standard come Keras LSTM. È l'opzione consigliata.
- Fornisci un'interfaccia nell'infrastruttura di conversione per implementazioni RNN definite dall'utente da collegare e convertire in LiteRT. Forniamo un paio di esempi predefiniti di questa conversione utilizzando le interfacce RNN LSTMCellSimple e LayerNormalizedLSTMCellSimple di Lingvo.
API Converter
La funzionalità fa parte della release di TensorFlow 2.3. È disponibile anche tramite tf-nightly pip o da head.
Questa funzionalità di conversione è disponibile quando si esegue la conversione in LiteRT tramite un SavedModel o direttamente dal modello Keras. Vedi esempi di utilizzo.
Da un 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()
Da 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 LiteRT Colab illustra l'utilizzo end-to-end con l'interprete LiteRT.
API RNN di TensorFlow supportate
Conversione LSTM di Keras (consigliata)
Supportiamo la conversione predefinita di Keras LSTM in LiteRT. Per informazioni dettagliate su come funziona, consulta l'interfaccia Keras LSTM e la logica di conversione qui.
È anche importante evidenziare il contratto LSTM di LiteRT rispetto alla definizione dell'operazione Keras:
- La dimensione 0 del tensore input è la dimensione del batch.
- La dimensione 0 del tensore recurrent_weight è il numero di output.
- I tensori weight e recurrent_kernel vengono trasposti.
- I tensori peso trasposto, kernel ricorrente trasposto e bias vengono divisi in 4 tensori di dimensioni uguali lungo la dimensione 0. Questi corrispondono a porta di input, porta di oblio, cella e porta di output.
Varianti LSTM di Keras
Time major
Gli utenti possono scegliere di utilizzare o meno la dimensione temporale. Keras LSTM aggiunge un attributo time-major negli attributi di definizione della funzione. Per LSTM sequenza unidirezionale, possiamo mappare semplicemente all'attributo time major di unidirecional_sequence_lstm.
LSTM bidirezionale
LSTM bidirezionale può essere implementato con due livelli LSTM di Keras, uno per l'avanti e uno per l'indietro. Vedi gli esempi qui. Una volta visualizzato l'attributo go_backward, lo riconosciamo come LSTM all'indietro, quindi raggruppiamo LSTM in avanti e all'indietro. Questo è il futuro del lavoro. Attualmente, questo crea due operazioni UnidirectionalSequenceLSTM nel modello LiteRT.
Esempi di conversione LSTM definita dall'utente
LiteRT fornisce anche un modo per convertire le implementazioni LSTM definite dall'utente. Qui utilizziamo l'LSTM di Lingvo come esempio di come può essere implementato. Per i dettagli, fai 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 nella relativa logica di conversione qui.
"Bring your own TensorFlow RNN" a LiteRT
Se l'interfaccia RNN di un utente è diversa da quelle standard supportate, esistono due opzioni:
Opzione 1: scrivi il codice dell'adattatore in TensorFlow Python per adattare l'interfaccia RNN all'interfaccia RNN di Keras. Ciò significa una tf.function con l'annotazione tf_implements sulla funzione dell'interfaccia RNN generata, identica a quella generata dal livello LSTM di Keras. Dopodiché, funzionerà la stessa API Conversion utilizzata per Keras LSTM.
Opzione 2:se quanto sopra non è possibile (ad es. la LSTM di Keras non ha alcune funzionalità attualmente esposte dall'operazione LSTM unita di LiteRT, come la normalizzazione dei livelli), estendi il convertitore LiteRT scrivendo codice di conversione personalizzato e collegalo al passaggio MLIR prepare-composite-functions qui. L'interfaccia della funzione deve essere trattata come un contratto API e deve contenere gli argomenti necessari per la conversione in operazioni LSTM LiteRT fuse, ad esempio input, bias, pesi, proiezione, normalizzazione del livello e così via. È preferibile che i tensori passati come argomenti a questa funzione abbiano un rango noto (ovvero RankedTensorType in MLIR). In questo modo, è molto più facile scrivere codice di conversione che può considerare questi tensori come RankedTensorType e aiutarli a trasformarli in tensori classificati corrispondenti agli operandi dell'operatore LiteRT unito.
Un esempio completo di questo flusso di conversione è la conversione da LSTMCellSimple a LiteRT di Lingvo.
LSTMCellSimple in Lingvo è definito qui. I modelli addestrati con questa cella LSTM possono essere convertiti in LiteRT nel seguente modo:
- Racchiudi tutti gli utilizzi di LSTMCellSimple in una tf.function con un'annotazione tf_implements etichettata come tale (ad es. lingvo.LSTMCellSimple sarebbe un buon nome di 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.
Estendi il passaggio prepare-composite-functions per inserire un'operazione composita personalizzata nella conversione dell'operazione LSTM fusa di LiteRT. Consulta il codice di conversione di LSTMCellSimple.
Il contratto di conversione:
I tensori peso e proiezione vengono trasposti.
L'{input, recurrent} a {cell, input gate, forget gate, output gate} vengono estratti sezionando il tensore dei pesi trasposto.
Il {bias} a {cell, input gate, forget gate, output gate} sono estratti sezionando il tensore di bias.
La proiezione viene estratta sezionando il tensore di proiezione trasposto.
La conversione simile è scritta per LayerNormalizedLSTMCellSimple.
Il resto dell'infrastruttura di conversione LiteRT, inclusi tutti i passaggi MLIR definiti e l'esportazione finale in LiteRT flatbuffer, può essere riutilizzato.
Problemi noti/limitazioni
- Al momento è supportata solo la conversione di LSTM Keras stateless (comportamento predefinito in Keras). La conversione di LSTM Keras stateful è un lavoro futuro.
- È comunque possibile modellare un livello LSTM di Keras stateful utilizzando il livello LSTM di Keras stateless sottostante e gestendo lo stato in modo esplicito nel programma utente. Un programma TensorFlow di questo tipo può comunque essere convertito in LiteRT utilizzando la funzionalità descritta qui.
- Al momento, LSTM bidirezionale è modellato come due operazioni UnidirectionalSequenceLSTM in LiteRT. Verrà sostituito con un singolo op BidirectionalSequenceLSTM.