Présentation
LiteRT permet de convertir les modèles TensorFlow RNN en opérations LSTM fusionnées de LiteRT. Les opérations fusionnées existent pour maximiser les performances de leurs implémentations de noyau sous-jacentes, ainsi que pour fournir une interface de niveau supérieur permettant de définir des transformations complexes telles que la quantification.
Comme il existe de nombreuses variantes d'API RNN dans TensorFlow, notre approche est double :
- Fournir une compatibilité native avec les API RNN TensorFlow standards, comme Keras LSTM. Il s'agit de l'option recommandée.
- Fournissez une interface dans l'infrastructure de conversion pour les implémentations RNN définies par l'utilisateur à brancher et à convertir en LiteRT. Nous fournissons quelques exemples prêts à l'emploi de ce type de conversion à l'aide des interfaces RNN LSTMCellSimple et LayerNormalizedLSTMCellSimple de Lingvo.
API Converter
Cette fonctionnalité fait partie de la version 2.3 de TensorFlow. Il est également disponible via le pip tf-nightly ou depuis la tête.
Cette fonctionnalité de conversion est disponible lorsque vous convertissez un SavedModel ou un modèle Keras directement en LiteRT. Consultez des exemples d'utilisation.
À partir d'un modèle enregistré
# 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()
À partir d'un modèle Keras
# build a Keras model
keras_model = build_keras_lstm(...)
# Convert the model.
converter = TFLiteConverter.from_keras_model(keras_model)
tflite_model = converter.convert()
Exemple
Le Colab Keras LSTM vers LiteRT illustre l'utilisation de bout en bout avec l'interpréteur LiteRT.
API TensorFlow RNN compatibles
Conversion Keras LSTM (recommandée)
Nous acceptons la conversion prête à l'emploi de Keras LSTM en LiteRT. Pour en savoir plus sur le fonctionnement de cette fonctionnalité, consultez l'interface Keras LSTM et la logique de conversion ici.
Il est également important de souligner le contrat LSTM de LiteRT par rapport à la définition de l'opération Keras :
- La dimension 0 du Tensor input correspond à la taille du lot.
- La dimension 0 du tenseur recurrent_weight correspond au nombre de sorties.
- Les Tensors weight et recurrent_kernel sont transposés.
- Les Tensors transposés weight, transposed_recurrent_kernel et bias sont divisés en quatre Tensors de taille égale le long de la dimension 0. Ils correspondent à la porte d'entrée, à la porte d'oubli, à la cellule et à la porte de sortie.
Variantes LSTM de Keras
Temps majeur
Les utilisateurs peuvent choisir d'utiliser ou non le format "time-major". Keras LSTM ajoute un attribut "time-major" dans les attributs de définition de fonction. Pour le LSTM de séquence unidirectionnelle, nous pouvons simplement mapper à l'attribut time major de unidirecional_sequence_lstm.
LSTM bidirectionnel
Les LSTM bidirectionnelles peuvent être implémentées avec deux couches LSTM Keras, une pour l'avant et une pour l'arrière. Pour voir des exemples, cliquez ici. Une fois que nous avons identifié l'attribut go_backward, nous le reconnaissons comme un LSTM inverse, puis nous regroupons les LSTM avant et inverse. Il s'agit d'un travail à venir. Actuellement, cela crée deux opérations UnidirectionalSequenceLSTM dans le modèle LiteRT.
Exemples de conversions LSTM définies par l'utilisateur
LiteRT permet également de convertir les implémentations LSTM définies par l'utilisateur. Ici, nous utilisons le LSTM de Lingvo comme exemple d'implémentation. Pour en savoir plus, veuillez consulter l'interface lingvo.LSTMCellSimple et la logique de conversion ici. Nous fournissons également un exemple pour une autre définition LSTM de Lingvo dans l'interface lingvo.LayerNormalizedLSTMCellSimple et sa logique de conversion ici.
"Bring your own TensorFlow RNN" vers LiteRT
Si l'interface RNN d'un utilisateur est différente de celles standards compatibles, plusieurs options s'offrent à lui :
Option 1 : Écrivez le code de l'adaptateur en Python TensorFlow pour adapter l'interface RNN à l'interface Keras RNN. Cela signifie qu'une tf.function avec annotation tf_implements sur la fonction de l'interface RNN générée est identique à celle générée par le calque LSTM Keras. Après cela, la même API de conversion utilisée pour Keras LSTM fonctionnera.
Option 2 : Si l'option ci-dessus n'est pas possible (par exemple, si le LSTM Keras manque de fonctionnalités actuellement exposées par l'opération LSTM fusionnée de LiteRT, comme la normalisation des couches), étendez le convertisseur LiteRT en écrivant un code de conversion personnalisé et en l'insérant dans le pass MLIR prepare-composite-functions ici. L'interface de la fonction doit être traitée comme un contrat d'API et doit contenir les arguments nécessaires à la conversion en opérations LSTM LiteRT fusionnées, c'est-à-dire l'entrée, le biais, les pondérations, la projection, la normalisation de couche, etc. Il est préférable que les Tensor transmis en tant qu'arguments à cette fonction aient un rang connu (c'est-à-dire RankedTensorType dans MLIR). Il est ainsi beaucoup plus facile d'écrire du code de conversion qui peut supposer que ces Tensors sont de type RankedTensorType et qui permet de les transformer en Tensors classés correspondant aux opérandes de l'opérateur LiteRT fusionné.
La conversion de LSTMCellSimple en LiteRT de Lingvo est un exemple complet de ce type de flux de conversion.
LSTMCellSimple dans Lingvo est défini ici. Les modèles entraînés avec cette cellule LSTM peuvent être convertis en LiteRT comme suit :
- Encapsulez toutes les utilisations de LSTMCellSimple dans une tf.function avec une annotation tf_implements libellée comme telle (par exemple, lingvo.LSTMCellSimple serait un bon nom d'annotation ici). Assurez-vous que la fonction tf.function générée correspond à l'interface de la fonction attendue dans le code de conversion. Il s'agit d'un contrat entre l'auteur du modèle qui ajoute l'annotation et le code de conversion.
Étendez le pass prepare-composite-functions pour brancher une opération composite personnalisée à la conversion d'opération LSTM fusionnée LiteRT. Consultez le code de conversion LSTMCellSimple.
Contrat de conversion :
Les Tensors weight et projection sont transposés.
Les {input, recurrent} vers {cell, input gate, forget gate, output gate} sont extraits en segmentant le Tensor de poids transposé.
Les {biais} vers {cellule, porte d'entrée, porte d'oubli, porte de sortie} sont extraits en segmentant le Tensor de biais.
La projection est extraite en segmentant le Tensor de projection transposé.
Une conversion similaire est écrite pour LayerNormalizedLSTMCellSimple.
Le reste de l'infrastructure de conversion LiteRT, y compris tous les passes MLIR définis, ainsi que l'exportation finale vers le flatbuffer LiteRT, peuvent être réutilisés.
Problèmes/Limites connus
- Actuellement, seule la conversion des LSTM Keras sans état (comportement par défaut dans Keras) est prise en charge. La conversion LSTM Keras avec état est une fonctionnalité à venir.
- Il est toujours possible de modéliser une couche LSTM Keras avec état à l'aide de la couche LSTM Keras sans état sous-jacente et de gérer l'état de manière explicite dans le programme utilisateur. Un tel programme TensorFlow peut toujours être converti en LiteRT à l'aide de la fonctionnalité décrite ici.
- Le LSTM bidirectionnel est actuellement modélisé sous la forme de deux opérations UnidirectionalSequenceLSTM dans LiteRT. Il sera remplacé par une seule opération BidirectionalSequenceLSTM.