Conversión de operadores de texto de TensorFlow a TensorFlow Lite

Ver en TensorFlow.org Ejecutar en Google Colab Ver en GitHub Descargar libreta

Descripción general

Los modelos de aprendizaje automático se implementan con frecuencia mediante TensorFlow Lite en dispositivos móviles, integrados y de IoT para mejorar la privacidad de los datos y reducir los tiempos de respuesta. Estos modelos a menudo requieren soporte para operaciones de procesamiento de texto. TensorFlow Text versión 2.7 y superior proporciona un rendimiento mejorado, tamaños binarios reducidos y operaciones específicamente optimizadas para su uso en estos entornos.

Operadores de texto

Las siguientes clases de texto de TensorFlow se pueden usar desde un modelo de TensorFlow Lite.

  • FastWordpieceTokenizer
  • WhitespaceTokenizer

Ejemplo de modelo

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

El siguiente ejemplo de código muestra el proceso de conversión y la interpretación en Python usando un modelo de prueba simple. Tenga en cuenta que la salida de un modelo no puede ser un objeto tf.RaggedTensor cuando usa TensorFlow Lite. Sin embargo, puede devolver los componentes de un objeto tf.RaggedTensor o convertirlo usando su función to_tensor . Consulte la guía RaggedTensor para obtener más detalles.

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)

Convierte el modelo TensorFlow a TensorFlow Lite

Al convertir un modelo de TensorFlow con operadores de texto de TensorFlow a TensorFlow Lite, debe indicarle al TFLiteConverter que hay operadores personalizados que usan el atributo allow_custom_ops como en el ejemplo a continuación. A continuación, puede ejecutar la conversión del modelo como lo haría normalmente. Revisa la documentación del convertidor TensorFlow Lite para obtener una guía detallada sobre los aspectos básicos de la conversión de modelos.

# 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

Inferencia

Para que el intérprete de TensorFlow Lite lea correctamente su modelo que contiene operadores de texto de TensorFlow, debe configurarlo para usar estos operadores personalizados y proporcionar métodos de registro para ellos. Use tf_text.tflite_registrar.SELECT_TFTEXT_OPS para proporcionar el conjunto completo de funciones de registro para los operadores de texto de TensorFlow admitidos para InterpreterWithCustomOps .

Tenga en cuenta que, si bien el siguiente ejemplo muestra la inferencia en Python, los pasos son similares en otros idiomas con algunas traducciones menores de API y la necesidad de construir el tflite_registrar en su binario. Consulte Inferencia de TensorFlow Lite para obtener más detalles.

# 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']} }

A continuación, se invoca el intérprete de TensorFlow Lite con la entrada, lo que proporciona un resultado que coincide con el resultado de TensorFlow anterior.

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']