Conversion des opérateurs de texte TensorFlow en TensorFlow Lite

Voir sur TensorFlow.org Exécuter dans Google Colab Afficher sur GitHub Télécharger le cahier

Aperçu

Les modèles d'apprentissage automatique sont fréquemment déployés à l'aide de TensorFlow Lite sur des appareils mobiles, intégrés et IoT afin d'améliorer la confidentialité des données et de réduire les temps de réponse. Ces modèles nécessitent souvent une prise en charge des opérations de traitement de texte. TensorFlow Text version 2.7 et ultérieure offre des performances améliorées, des tailles binaires réduites et des opérations spécifiquement optimisées pour une utilisation dans ces environnements.

Opérateurs de texte

Les classes TensorFlow Text suivantes peuvent être utilisées à partir d'un modèle TensorFlow Lite.

  • FastWordpieceTokenizer
  • WhitespaceTokenizer

Exemple de modèle

pip install -U tensorflow-text
from absl import app
import numpy as np
import tensorflow as tf
import tensorflow_text as tf_text

from tensorflow.lite.python import interpreter

L'exemple de code suivant montre le processus de conversion et l'interprétation en Python à l'aide d'un modèle de test simple. Notez que la sortie d'un modèle ne peut pas être un objet tf.RaggedTensor lorsque vous utilisez TensorFlow Lite. Cependant, vous pouvez renvoyer les composants d'un objet tf.RaggedTensor ou le convertir à l'aide de sa fonction to_tensor . Voir le guide RaggedTensor pour plus de détails.

class TokenizerModel(tf.keras.Model):

  def __init__(self, **kwargs):
    super().__init__(**kwargs)
    self.tokenizer = tf_text.WhitespaceTokenizer()

  @tf.function(input_signature=[
      tf.TensorSpec(shape=[None], dtype=tf.string, name='input')
  ])
  def call(self, input_tensor):
    return { 'tokens': self.tokenizer.tokenize(input_tensor).flat_values }
# Test input data.
input_data = np.array(['Some minds are better kept apart'])

# Define a Keras model.
model = TokenizerModel()

# Perform TensorFlow Text inference.
tf_result = model(tf.constant(input_data))
print('TensorFlow result = ', tf_result['tokens'])
TensorFlow result =  tf.Tensor([b'Some' b'minds' b'are' b'better' b'kept' b'apart'], shape=(6,), dtype=string)

Convertir le modèle TensorFlow en TensorFlow Lite

Lors de la conversion d'un modèle TensorFlow avec des opérateurs TensorFlow Text vers TensorFlow Lite, vous devez indiquer au TFLiteConverter qu'il existe des opérateurs personnalisés utilisant l'attribut allow_custom_ops comme dans l'exemple ci-dessous. Vous pouvez ensuite exécuter la conversion de modèle comme vous le feriez normalement. Consultez la documentation du convertisseur TensorFlow Lite pour un guide détaillé sur les bases de la conversion de modèle.

# Convert to TensorFlow Lite.
converter = tf.lite.TFLiteConverter.from_keras_model(model)
converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS]
converter.allow_custom_ops = True
tflite_model = converter.convert()
2022-02-01 12:09:02.062677: W tensorflow/python/util/util.cc:368] Sets are not currently considered sequences, but this may change in the future, so consider avoiding using them.
INFO:tensorflow:Assets written to: /tmp/tmpiiuhjdn6/assets
2022-02-01 12:09:03.705144: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:363] Ignored output_format.
WARNING:absl:Buffer deduplication procedure will be skipped when flatbuffer library is not properly loaded
2022-02-01 12:09:03.705185: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:366] Ignored drop_control_dependency.
2022-02-01 12:09:03.921830: W tensorflow/compiler/mlir/lite/flatbuffer_export.cc:1902] The following operation(s) need TFLite custom op implementation(s):
Custom ops: TFText>WhitespaceTokenizeWithOffsetsV2
Details:
    tf.TFText>WhitespaceTokenizeWithOffsetsV2(tensor<?x!tf_type.string>, tensor<!tf_type.string>) -> (tensor<?x!tf_type.string>, tensor<?xi64>, tensor<?xi32>, tensor<?xi32>) : {device = ""}
See instructions: https://www.tensorflow.org/lite/guide/ops_custom

Inférence

Pour que l'interpréteur TensorFlow Lite lise correctement votre modèle contenant des opérateurs de texte TensorFlow, vous devez le configurer pour utiliser ces opérateurs personnalisés et leur fournir des méthodes d'enregistrement. Utilisez tf_text.tflite_registrar.SELECT_TFTEXT_OPS pour fournir la suite complète de fonctions d'enregistrement pour les opérateurs de texte TensorFlow pris en charge à InterpreterWithCustomOps .

Notez que bien que l'exemple ci-dessous montre l'inférence en Python, les étapes sont similaires dans d'autres langages avec quelques traductions mineures de l'API et la nécessité de construire le tflite_registrar dans votre binaire. Voir Inférence TensorFlow Lite pour plus de détails.

# Perform TensorFlow Lite inference.
interp = interpreter.InterpreterWithCustomOps(
    model_content=tflite_model,
    custom_op_registerers=tf_text.tflite_registrar.SELECT_TFTEXT_OPS)
interp.get_signature_list()
{'serving_default': {'inputs': ['input'], 'outputs': ['tokens']} }

Ensuite, l'interpréteur TensorFlow Lite est appelé avec l'entrée, fournissant un résultat qui correspond au résultat TensorFlow ci-dessus.

tokenize = interp.get_signature_runner('serving_default')
output = tokenize(input=input_data)
print('TensorFlow Lite result = ', output['tokens'])
TensorFlow Lite result =  [b'Some' b'minds' b'are' b'better' b'kept' b'apart']