将 TensorFlow Text 运算符转换为 TensorFlow Lite

在 TensorFlow.org 上查看

在 Google Colab 中运行

在 GitHub 上查看

下载笔记本

概述

机器学习模型通常使用 TensorFlow Lite 部署到移动设备、嵌入式设备和物联网设备，以提高数据隐私并降低响应时间。这些模型通常需要支持文本处理操作。TensorFlow Text 2.7 及更高版本提供了改进的性能、更小的二进制文件大小以及专门针对这些环境优化的操作。

文本运算符

以下 TensorFlow Text 类和函数可以在 TensorFlow Lite 模型中使用。

模型示例

pip install -U "tensorflow-text==2.11.*"

from absl import app
import numpy as np
import tensorflow as tf
import tensorflow_text as tf_text

from tensorflow.lite.python import interpreter

2024-06-25 11:43:48.257887: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-06-25 11:43:48.922853: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer.so.7'; dlerror: libnvinfer.so.7: cannot open shared object file: No such file or directory
2024-06-25 11:43:48.922950: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libnvinfer_plugin.so.7'; dlerror: libnvinfer_plugin.so.7: cannot open shared object file: No such file or directory
2024-06-25 11:43:48.922960: W tensorflow/compiler/tf2tensorrt/utils/py_utils.cc:38] TF-TRT Warning: Cannot dlopen some TensorRT libraries. If you would like to use Nvidia GPU with TensorRT, please make sure the missing libraries mentioned above are installed properly.

以下代码示例展示了使用简单测试模型在 Python 中的转换过程和解释。请注意，当您使用 TensorFlow Lite 时，模型的输出不能是 tf.RaggedTensor 对象。但是，您可以返回 tf.RaggedTensor 对象的组件，或使用其 to_tensor 函数进行转换。有关更多详细信息，请参阅 RaggedTensor 指南。

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

2024-06-25 11:43:50.638719: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2024-06-25 11:43:50.638821: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublas.so.11'; dlerror: libcublas.so.11: cannot open shared object file: No such file or directory
2024-06-25 11:43:50.638883: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcublasLt.so.11'; dlerror: libcublasLt.so.11: cannot open shared object file: No such file or directory
2024-06-25 11:43:50.638942: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcufft.so.10'; dlerror: libcufft.so.10: cannot open shared object file: No such file or directory
2024-06-25 11:43:50.693816: W tensorflow/compiler/xla/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcusparse.so.11'; dlerror: libcusparse.so.11: cannot open shared object file: No such file or directory
2024-06-25 11:43:50.694015: W tensorflow/core/common_runtime/gpu/gpu_device.cc:1934] Cannot dlopen some GPU libraries. Please make sure the missing libraries mentioned above are installed properly if you would like to use GPU. Follow the guide at https://tensorflowcn.cn/install/gpu for how to download and setup the required libraries for your platform.
Skipping registering GPU devices...
TensorFlow result =  tf.Tensor([b'Some' b'minds' b'are' b'better' b'kept' b'apart'], shape=(6,), dtype=string)

将 TensorFlow 模型转换为 TensorFlow Lite

将包含 TensorFlow Text 运算符的 TensorFlow 模型转换为 TensorFlow Lite 时，您需要使用 allow_custom_ops 属性向 TFLiteConverter 指示存在自定义运算符，如以下示例所示。然后，您可以像往常一样运行模型转换。有关模型转换基础知识的详细指南，请查看 TensorFlow Lite 转换器文档。

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

INFO:tensorflow:Assets written to: /tmpfs/tmp/tmpxukrx8mh/assets
2024-06-25 11:43:53.298066: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:362] Ignored output_format.
2024-06-25 11:43:53.298105: W tensorflow/compiler/mlir/lite/python/tf_tfl_flatbuffer_helpers.cc:365] Ignored drop_control_dependency.
2024-06-25 11:43:53.495164: W tensorflow/compiler/mlir/lite/flatbuffer_export.cc:2057] 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://tensorflowcn.cn/lite/guide/ops_custom

推理

为了使 TensorFlow Lite 解释器能够正确读取包含 TensorFlow Text 运算符的模型，您必须配置它以使用这些自定义运算符，并为它们提供注册方法。使用 tf_text.tflite_registrar.SELECT_TFTEXT_OPS 为 InterpreterWithCustomOps 提供支持的 TensorFlow Text 运算符的完整注册函数集。

请注意，虽然以下示例展示了在 Python 中进行推理，但这些步骤在其他语言中也类似，只是 API 翻译略有不同，并且需要将 tflite_registrar 构建到您的二进制文件中。有关更多详细信息，请参阅 TensorFlow Lite 推理。

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

接下来，使用输入调用 TensorFlow Lite 解释器，提供与上述 TensorFlow 结果匹配的结果。

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