Como converter operadores de texto do TensorFlow em TensorFlow Lite

Veja no TensorFlow.org Executar no Google Colab Ver no GitHub Baixar caderno

Visão geral

Os modelos de aprendizado de máquina são frequentemente implantados usando o TensorFlow Lite em dispositivos móveis, incorporados e IoT para melhorar a privacidade dos dados e reduzir os tempos de resposta. Esses modelos geralmente exigem suporte para operações de processamento de texto. O TensorFlow Text versão 2.7 e superior oferece desempenho aprimorado, tamanhos binários reduzidos e operações especificamente otimizadas para uso nesses ambientes.

Operadores de texto

As seguintes classes de texto do TensorFlow podem ser usadas em um modelo do TensorFlow Lite.

  • FastWordpieceTokenizer
  • WhitespaceTokenizer

Exemplo 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

O exemplo de código a seguir mostra o processo de conversão e interpretação em Python usando um modelo de teste simples. Observe que a saída de um modelo não pode ser um objeto tf.RaggedTensor quando você estiver usando o TensorFlow Lite. No entanto, você pode retornar os componentes de um objeto tf.RaggedTensor ou convertê-lo usando sua função to_tensor . Consulte o guia RaggedTensor para obter mais detalhes.

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)

Converter o modelo do TensorFlow para o TensorFlow Lite

Ao converter um modelo do TensorFlow com operadores de texto do TensorFlow para o TensorFlow Lite, você precisa indicar ao TFLiteConverter que há operadores personalizados usando o atributo allow_custom_ops como no exemplo abaixo. Você pode então executar a conversão do modelo como faria normalmente. Revise a documentação do conversor do TensorFlow Lite para obter um guia detalhado sobre os fundamentos da conversão 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

Inferência

Para que o interpretador do TensorFlow Lite leia corretamente seu modelo contendo operadores de texto do TensorFlow, você deve configurá-lo para usar esses operadores personalizados e fornecer métodos de registro para eles. Use tf_text.tflite_registrar.SELECT_TFTEXT_OPS para fornecer o conjunto completo de funções de registro para os operadores de texto do TensorFlow compatíveis para InterpreterWithCustomOps .

Observe que, embora o exemplo abaixo mostre inferência em Python, as etapas são semelhantes em outros idiomas com algumas traduções de API menores e a necessidade de criar o tflite_registrar em seu binário. Consulte Inferência do TensorFlow Lite para obter mais detalhes.

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

Em seguida, o interpretador do TensorFlow Lite é invocado com a entrada, fornecendo um resultado que corresponde ao resultado do TensorFlow acima.

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