將 TensorFlow Text 運算元轉換為 TensorFlow Lite

在 TensorFlow.org 上檢視 在 Google Colab 中執行 在 GitHub 上檢視 下載筆記本

總覽

機器學習模型經常使用 TensorFlow Lite 部署到行動裝置、嵌入式裝置和 IoT 裝置，以改善資料隱私並縮短回應時間。這些模型通常需要支援文字處理作業。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://tensorflow.dev.org.tw/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 時，您需要告知 TFLiteConverter 有自訂運算元，方法是使用 allow_custom_ops 屬性，如下列範例所示。然後，您可以像平常一樣執行模型轉換。如需模型轉換基本概念的詳細指南，請參閱 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://tensorflow.dev.org.tw/lite/guide/ops_custom

推論

為了讓 TensorFlow Lite 解譯器正確讀取包含 TensorFlow Text 運算元的模型，您必須將其設定為使用這些自訂運算元，並為其提供註冊方法。使用 tf_text.tflite_registrar.SELECT_TFTEXT_OPS 為支援的 TensorFlow Text 運算元完整套件的註冊函式提供給 InterpreterWithCustomOps。

請注意，雖然下列範例顯示 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']