Conversion de TensorFlow RNN en TensorFlow Lite

Aperçu

TensorFlow Lite prend en charge la conversion des modèles TensorFlow RNN en opérations LSTM fusionnées de TensorFlow Lite. 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 pour définir des transformations complexes telles que la quantification.

Puisqu'il existe de nombreuses variantes d'API RNN dans TensorFlow, notre approche a été double :

  1. Fournissez une prise en charge native des API TensorFlow RNN standard telles que Keras LSTM. C'est l'option recommandée.
  2. Fournissez une interface dans l'infrastructure de conversion pour que les implémentations RNN définies par l'utilisateur puissent se connecter et être converties en TensorFlow Lite. Nous fournissons quelques exemples prêts à l'emploi d'une telle conversion à l'aide des interfaces RNN LSTMCellSimple et LayerNormalizedLSTMCellSimple de Lingvo.

API du convertisseur

La fonctionnalité fait partie de la version TensorFlow 2.3. Il est également disponible via le pip tf-nightly ou depuis head.

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

Du 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 RNN TensorFlow prises en charge

Nous prenons en charge la conversion prête à l'emploi de Keras LSTM vers TensorFlow Lite. Pour plus de détails sur la façon dont cela fonctionne, veuillez vous référer à l' interface Keras LSTM et à la logique de conversion ici .

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

  1. La dimension 0 du tenseur d'entrée est la taille du lot.
  2. La dimension 0 du tenseur recurrent_weight est le nombre de sorties.
  3. Les tenseurs de poids et de recurrent_kernel sont transposés.
  4. Les tenseurs de poids transposé, de recurrent_kernel et de biais transposés sont divisés en 4 tenseurs de taille égale le long de la dimension 0. Ceux-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 entre temps majeur ou pas de temps majeur. Keras LSTM ajoute un attribut time-major dans les attributs def de la fonction. Pour la séquence unidirectionnelle LSTM, nous pouvons simplement mapper sur l'attribut time major d'unidirecional_sequence_lstm.

LSTM bidirectionnel

Le LSTM bidirectionnel peut être implémenté avec deux couches Keras LSTM, une pour l'avant et une pour l'arrière, voir les exemples ici . Une fois que nous voyons l'attribut go_backward, nous le reconnaissons comme un LSTM arrière, puis nous regroupons les LSTM avant et arrière. C'est un travail futur. Actuellement, cela crée deux opérations UnidirectionnelleSequenceLSTM dans le modèle TensorFlow Lite.

Exemples de conversion LSTM définis par l'utilisateur

TensorFlow Lite fournit également un moyen de convertir les implémentations LSTM définies par l'utilisateur. Ici, nous utilisons le LSTM de Lingvo comme exemple de la façon dont cela peut être mis en œuvre. Pour plus de détails, veuillez vous référer à l' interface lingvo.LSTMCellSimple et à la logique de conversion ici . Nous fournissons également un exemple d'une autre définition LSTM de Lingvo dans l'interface lingvo.LayerNormalizedLSTMCellSimple et sa logique de conversion ici .

« Apportez votre propre RNN TensorFlow » à TensorFlow Lite

Si l'interface RNN d'un utilisateur est différente de celles standard prises en charge, il existe plusieurs options :

Option 1 : Écrivez le code de l'adaptateur dans TensorFlow python pour adapter l'interface RNN à l'interface Keras RNN. Cela signifie un tf.function avec une annotation tf_implements sur la fonction de l'interface RNN générée qui est identique à celle générée par la couche Keras LSTM. Après cela, la même API de conversion utilisée pour Keras LSTM fonctionnera.

Option 2 : si ce qui précède n'est pas possible (par exemple, il manque au Keras LSTM certaines fonctionnalités actuellement exposées par l'opération LSTM fusionnée de TensorFlow Lite, comme la normalisation des couches), étendez le convertisseur TensorFlow Lite en écrivant un code de conversion personnalisé et branchez-le dans le fichier de préparation. -fonctions composites MLIR-passez ici . L'interface de la fonction doit être traitée comme un contrat API et doit contenir les arguments nécessaires à la conversion en opérations TensorFlow Lite LSTM fusionnées - c'est-à-dire entrée, biais, poids, projection, normalisation de couche, etc. Il est préférable que les tenseurs soient passés en arguments à cette fonction pour avoir un rang connu (c'est-à-dire RankedTensorType dans MLIR). Cela rend beaucoup plus facile l'écriture de code de conversion qui peut prendre ces tenseurs comme RankedTensorType et aide à les transformer en tenseurs classés correspondant aux opérandes de l'opérateur TensorFlow Lite fusionné.

Un exemple complet d'un tel flux de conversion est la conversion LSTMCellSimple vers TensorFlow Lite de Lingvo.

Le LSTMCellSimple dans Lingvo est défini ici . Les modèles entraînés avec cette cellule LSTM peuvent être convertis en TensorFlow Lite comme suit :

  1. Enveloppez toutes les utilisations de LSTMCellSimple dans un tf.function avec une annotation tf_implements étiquetée comme telle (par exemple, lingvo.LSTMCellSimple serait un bon nom d'annotation ici). Assurez-vous que le tf.function généré correspond à l'interface de la fonction attendue dans le code de conversion. Il s'agit d'un contrat entre l'auteur du modèle ajoutant l'annotation et le code de conversion.
  2. Étendez l'étape de préparation-composite-fonctions pour connecter une opération composite personnalisée à la conversion d'opération LSTM fusionnée TensorFlow Lite. Voir le code de conversion LSTMCellSimple .

    Le contrat de transformation :

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

  4. Les {entrée, récurrentes} à {cellule, porte d'entrée, porte d'oubli, porte de sortie} sont extraites en découpant le tenseur de poids transposé.

  5. Les {bias} à {cellule, porte d'entrée, porte d'oubli, porte de sortie} sont extraits en découpant le tenseur de biais.

  6. La projection est extraite en découpant le tenseur 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, peuvent être réutilisés.

Problèmes/limites connus

  1. Actuellement, la conversion de Keras LSTM sans état est uniquement prise en charge (comportement par défaut dans Keras). La conversion Stateful Keras LSTM est un travail futur.
  2. Il est toujours possible de modéliser une couche Keras LSTM avec état en utilisant la couche Keras LSTM sans état sous-jacente et en gérant l'état explicitement 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 LSTM bidirectionnel est actuellement modélisé comme deux opérations UnidirectionnelSequenceLSTM dans TensorFlow Lite. Ceci sera remplacé par une seule opération BidirectionnelleSequenceLSTM.