Conversion de RNN TensorFlow en TensorFlow Lite

Présentation

TensorFlow Lite permet de convertir des modèles de RNN TensorFlow en opérations LSTM fusionnées de TensorFlow Lite. Les opérations fusionnées existent pour maximiser les performances des implémentations de noyau sous-jacentes, et fournir une interface de niveau supérieur pour définir des transformations complexes telles que la quantification.

Étant donné qu'il existe de nombreuses variantes des API RNN dans TensorFlow, notre approche a été doublée:

  1. Bénéficiez d'une compatibilité native avec les API RNN TensorFlow standards telles que Keras LSTM. Il s'agit de l'option recommandée.
  2. Fournissez une interface dans l'infrastructure de conversion pour les implémentations de RNN définies par l'utilisateur définies par l'utilisateur afin de les connecter et de les convertir dans TensorFlow Lite. 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 tf-nightly ou depuis la tête.

Cette fonctionnalité de conversion est disponible lors de la conversion vers TensorFlow Lite via un SavedModel ou directement à partir du modèle Keras. Voir des exemples d'utilisation

À partir du 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

Keras LSTM vers TensorFlow Lite Colab illustre l'utilisation de bout en bout avec l'interpréteur TensorFlow Lite.

API TensorFlow RNN compatibles

La solution prête à l'emploi permet la conversion de Keras LSTM en TensorFlow Lite. Pour en savoir plus, consultez l'interface Keras LSTM et la logique de conversion.

Il est également important de mettre en évidence le contrat LSTM de TensorFlow Lite par rapport à la définition des opérations Keras:

  1. La dimension 0 du Tensor d'entrée correspond à la taille de lot.
  2. La dimension 0 du Tensor recurrent_weight correspond au nombre de sorties.
  3. Les Tensors weight et recurrent_kernel sont transposés.
  4. Les Tensors "weight transposé", "recurrent_kernel" et "bias" sont divisés en quatre Tensors de taille égale le long de la dimension 0. Celles-ci correspondent à la porte d'entrée, la porte d'oubli, la cellule et la porte de sortie.

Variantes Keras LSTM

Temps majeur

Les utilisateurs peuvent choisir une date de grandeur ou aucune. Keras LSTM ajoute un attribut temporel dans les attributs de définition de la fonction. Pour la séquence unidirectionnelle LSTM, il suffit de la mapper à l'attribut temporel majeur d'unidirecional_Sequence_lstm.

BiDirectional LSTM

Le modèle LSTM bidirectionnel peut être mis en œuvre avec deux couches Keras LSTM, une pour l'avant et une pour l'arrière. Pour consulter des exemples, cliquez ici. Une fois que nous avons détecté l'attribut "go_backward", nous le reconnaissons en tant que LSTM reculé, puis nous regroupons LSTM avant et arrière. Il s'agit d'un travail futur. Actuellement, deux opérations UnidirectionalSequenceLSTM sont créées dans le modèle TensorFlow Lite.

Exemples de conversions LSTM définies par l'utilisateur

TensorFlow Lite permet également de convertir des implémentations LSTM définies par l'utilisateur. Ici, nous utilisons LSTM de Lingvo comme exemple de mise en œuvre. Pour en savoir plus, consultez l'interface lingvo.LSTMCellSimple et la logique de conversion sur cette page. Nous fournissons également un exemple pour une autre définition LSTM de Lingvo dans l'interface lingvo.LayerNormalizedLSTMCellSimple et pour sa logique de conversion ici.

"Bring Your Own TensorFlow RNN" dans TensorFlow Lite

Si l'interface RNN d'un utilisateur est différente des interfaces standards prises en charge, deux options s'offrent à vous:

Option 1:écrire le code d'adaptateur en Python TensorFlow pour adapter l'interface RNN à l'interface Keras RNN. Cela signifie que la fonction tf.function avec l'annotation tf_implements est identique à celle générée par la couche Keras LSTM sur la fonction de l'interface RNN générée. La même API de conversion que celle utilisée pour Keras LSTM fonctionnera ensuite.

Option 2:si ce n'est pas possible (par exemple, si Keras LSTM n'a pas accès à une fonctionnalité actuellement exposée par la normalisation de couches op fusée LSTM de TensorFlow Lite), étendez le convertisseur TensorFlow Lite en écrivant un code de conversion personnalisé et en le connectant à l'outil MLIR-pass de préparation des fonctions composites ici. L'interface de la fonction doit être traitée comme un contrat d'API et contenir les arguments nécessaires à la conversion en opérations TensorFlow Lite LSTM fusionnées (entrée, biais, pondérations, projection, normalisation des couches, etc.). Il est préférable que les Tensors transmis en tant qu'arguments à cette fonction aient un rang connu (par exemple, RankedTensorType dans MLIR). Il est ainsi beaucoup plus facile d'écrire un code de conversion pouvant supposer que ces Tensors sont classés comme RankedTensorType et les transforme en Tensors classés correspondant aux opérandes de l'opérateur TensorFlow Lite fusionnés.

La conversion LSTMCellSimple de Lingvo en TensorFlow Lite est un exemple complet de ce type de flux de conversion.

La méthode LSTMCellSimple en Lingvo est définie ici. Les modèles entraînés avec cette cellule LSTM peuvent être convertis au format TensorFlow Lite comme suit:

  1. Encapsulez toutes les utilisations de LSTMCellSimple dans une fonction tf.function avec une annotation tf_Implements libellée comme telle (par exemple, lingvo.LSTMCellSimple serait un bon nom d'annotation dans ce cas). 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.
  2. Étendez la transmission des fonctions composites pour connecter une opération composite personnalisée à la conversion d'opération Fused LSTM de TensorFlow Lite. Consultez le code de conversion LSTMCellSimple.

    Le contrat de conversion:

  3. Les Tensors de poids et de projection sont transposés.

  4. Les valeurs {input, recurrent} à {cell, input gate, forgot gate, output gate} sont extraites en découpant le Tensor de pondération transposé.

  5. Le {bias} à {cell, input gate, forgot gate, output gate} est extrait en découpant le Tensor de biais.

  6. La projection est extraite en coupant le Tensor de projection transposé.

  7. Une conversion similaire est écrite pour LayerNormalizedLSTMCellSimple.

  8. Le reste de l'infrastructure de conversion TensorFlow Lite, y compris toutes les passes MLIR définies ainsi que l'exportation finale vers le tampon plat TensorFlow Lite, peut être réutilisé.

Limites et problèmes connus

  1. Actuellement, cette fonctionnalité n'est compatible qu'avec la conversion de Keras LSTM sans état (comportement par défaut dans Keras). La conversion avec état Keras LSTM est une tâche future.
  2. Il est toujours possible de modéliser une couche Keras LSTM avec état à l'aide de la couche Keras LSTM sans état sous-jacente et de gérer explicitement l'état dans le programme utilisateur. Un tel programme TensorFlow peut toujours être converti en TensorFlow Lite à l'aide de la fonctionnalité décrite ici.
  3. Le modèle LSTM bidirectionnel est actuellement modélisé comme deux opérations UnidirectionalSequenceLSTM dans TensorFlow Lite. Elle sera remplacée par une seule opération BidirectionalSequenceLSTM.