Conversion de RNN TensorFlow en LiteRT

Présentation

LiteRT permet de convertir des modèles TensorFlow RNN en modèles LiteRT les opérations LSTM fusionnées. Les opérations fusionnées permettent d'optimiser les performances leurs implémentations de noyau sous-jacentes, et fournissent un niveau pour définir des transformations complexes comme la quantification.

Comme il existe de nombreuses variantes des API RNN dans TensorFlow, notre approche a été en deux temps:

  1. Fournir une compatibilité native avec les API TensorFlow RNN standards comme Keras LSTM Il s'agit de l'option recommandée.
  2. fournir une interface vers l'infrastructure de conversion pour mises en œuvre de RNN définies par l'utilisateur à brancher et à convertir en LiteRT. Voici quelques exemples prêts à l'emploi de conversion à l'aide de LSTMCellSimple et LayerNormalizedLSTMCellSimple d'interfaces RNN.

API Converter

Cette fonctionnalité fait partie de la version 2.3 de TensorFlow. Vous pouvez également y accéder via le tf-nightly ou "head".

Cette fonctionnalité de conversion est disponible lors de la conversion au format LiteRT. via un SavedModel ou directement depuis le modèle Keras. Consultez 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 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 LiteRT Colab illustre l'utilisation de bout en bout avec l'interpréteur LiteRT.

API des RNN TensorFlow compatibles

Nous acceptons la conversion prête à l'emploi de Keras LSTM en LiteRT. Pour Pour en savoir plus, consultez le Interface Keras LSTM et à la logique de conversion cliquez ici.

Il est également important de souligner que le contrat LSTM de LiteRT respecte à la définition de l'opération Keras:

  1. La dimension 0 du Tensor input correspond à la taille de lot.
  2. La dimension 0 du Tensor recurrent_weight est le nombre de de sortie.
  3. Les Tensors weight et recurrent_kernel sont transposés.
  4. Les Tensors "poids transposé", "recurrent_kernel" transposés et "bias" sont en 4 Tensors de taille égale le long de la dimension 0. Elles correspondent à porte d'entrée, porte de suppression, porte de cellule et porte de sortie.

Variantes Keras LSTM

Temps important

Les utilisateurs peuvent choisir "Très grand" ou non. Keras LSTM ajoute une couche de dans les attributs de définition de la fonction. Pour la séquence unidirectionnelle LSTM, il suffit de mapper vers le paramètre attribut "time" [temps] principal.

BiDirectional LSTM

La méthode LSTM bidirectionnelle peut être implémentée avec deux couches Keras LSTM, l'une pour vers l'avant et l'autre vers l'arrière. Voir des exemples cliquez ici. Une fois que nous avons vu l'attribut go_backward, nous le reconnaissons comme LSTM ascendant, puis nous regroupons en arrière-plan. Il s'agit d'un travail futur. Actuellement, crée deux opérations UnidirectionalSequenceLSTM dans le fichier LiteRT du modèle.

Exemples de conversions LSTM définies par l'utilisateur

LiteRT permet également de convertir les requêtes LSTM définies par l'utilisateur. mises en œuvre. Ici, nous utilisons la technique LSTM de Lingvo pour illustrer mise en œuvre. Pour en savoir plus, consultez les Interface lingvo.LSTMCellSimple et la logique de conversion cliquez ici. Nous fournissons également un exemple pour une autre des définitions LSTM de Lingvo dans Interface lingvo.LayerNormalizedLSTMCellSimple et sa logique de conversion cliquez ici.

"Bring Your Own TensorFlow RNN" sur LiteRT

Si l'interface RNN d'un utilisateur est différente de celles standard prises en charge, plusieurs options s'offrent à vous:

Option 1:écrire le code de l'adaptateur dans TensorFlow Python pour adapter l'interface RNN à l'interface Keras RNN. Cela signifie qu'une fonction tf.function avec tf_implements annotation sur la fonction de l'interface RNN générée, qui est identique à celle générée par la couche Keras LSTM. Ensuite, la même API de conversion que celle utilisée pour Keras LSTM fonctionne correctement.

Option 2:Si ce qui précède n'est pas possible (par exemple, si le pipeline LSTM Keras n'inclut pas certaines actuellement exposée par l'opération Fused LSTM de LiteRT, comme normalisation des couches), puis étendez le convertisseur LiteRT en écrivant code de conversion personnalisé et l'insérer dans la fonction MLIR-pass cliquez ici. L'interface de la fonction doit être traitée comme un contrat d'API et doit contenant les arguments nécessaires à la conversion en Fused LiteRT LSTM opérations (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 (RankedTensorType dans MLIR). Cela facilite grandement la rédaction code de conversion qui peut supposer que ces Tensors sont "RankedTensorType" et aide les transformer en Tensors classés correspondant à l'ensemble de données LiteRT fusionné les opérandes de l'opérateur.

Le modèle LSTMCellSimple de Lingvo, qui permet Conversion LiteRT.

LSTMCellSimple dans Lingvo est défini cliquez ici. Les modèles entraînés avec cette cellule LSTM peuvent être convertis au format LiteRT en tant que ce qui suit:

  1. Encapsulez toutes les utilisations de LSTMCellSimple dans tf.function avec tf_implements une annotation libellée comme telle (par exemple, lingvo.LSTMCellSimple serait nom d'annotation approprié ici). Assurez-vous que la fonction tf.function générée correspond à l'interface de la fonction attendue dans le code de conversion. Ce est un contrat entre l'auteur du modèle qui ajoute l'annotation et code de conversion.
  2. Étendre la passe prepare-composite-functions pour insérer une opération composite personnalisée la conversion d'opérations LSTM fusionnées LiteRT. Voir LSTMCellSimple code de conversion.

    Le contrat de conversion:

  3. Les Tensors Pondération et projection sont transposés.

  4. Le paramètre {input, recurrent} à {cell, input gate, delete gate, output (sortie) gate} sont extraits en segmentant le Tensor de pondération transposée.

  5. La fonction {bias} à {bias} est en segmentant le Tensor de biais.

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

  7. Une conversion similaire est écrite pour LayerNormalizedLSTMCellSimple.

  8. Le reste de l'infrastructure de conversion LiteRT, y compris toutes les Pass MLIR et l'exportation finale vers le format LiteRT Flatbuffer peut être réutilisés.

Limites et problèmes connus

  1. Actuellement, seule la conversion de Keras LSTM sans état (par défaut, dans Keras). La conversion LSTM Keras avec état est un travail futur.
  2. Il est toujours possible de modéliser une couche Keras LSTM avec état à l'aide de la méthode la couche Keras LSTM sous-jacente sans état et gérer explicitement l'état dans le programme utilisateur. Un tel programme TensorFlow peut toujours être converti LiteRT utilisant la fonctionnalité décrite ici.
  3. Le modèle LSTM bidirectionnel est actuellement modélisé sous la forme de deux blocs UnidirectionalSequenceLSTM dans LiteRT. Il sera remplacé par un seul Opération BidirectionalSequenceLSTM