Chuyển đổi toán tử TensorFlow Text thành TensorFlow Lite

Xem trên TensorFlow.org Chạy trong Google Colab Xem trên GitHub Tải xuống sổ ghi chép

Tổng quat

Các mô hình học máy thường xuyên được triển khai bằng cách sử dụng TensorFlow Lite cho các thiết bị di động, nhúng và IoT để cải thiện quyền riêng tư của dữ liệu và giảm thời gian phản hồi. Các mô hình này thường yêu cầu hỗ trợ các hoạt động xử lý văn bản. TensorFlow Text phiên bản 2.7 trở lên cung cấp hiệu suất được cải thiện, giảm kích thước nhị phân và các hoạt động được tối ưu hóa đặc biệt để sử dụng trong các môi trường này.

Toán tử văn bản

Các lớp TensorFlow Text sau đây có thể được sử dụng từ bên trong mô hình TensorFlow Lite.

  • FastWordpieceTokenizer
  • WhitespaceTokenizer

Ví dụ về mô hình

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

Ví dụ mã sau đây cho thấy quá trình chuyển đổi và diễn giải bằng Python bằng cách sử dụng một mô hình thử nghiệm đơn giản. Lưu ý rằng đầu ra của một mô hình không được là đối tượng tf.RaggedTensor khi bạn đang sử dụng TensorFlow Lite. Tuy nhiên, bạn có thể trả về các thành phần của đối tượng tf.RaggedTensor hoặc chuyển đổi nó bằng cách sử dụng hàm to_tensor của nó. Xem hướng dẫn RaggedTensor để biết thêm chi tiết.

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)

Chuyển đổi mô hình TensorFlow thành TensorFlow Lite

Khi chuyển đổi mô hình TensorFlow với các toán tử TensorFlow Text thành TensorFlow Lite, bạn cần chỉ ra cho TFLiteConverter rằng có các toán tử tùy chỉnh sử dụng thuộc tính allow_custom_ops như trong ví dụ bên dưới. Sau đó, bạn có thể chạy chuyển đổi mô hình như bình thường. Xem lại tài liệu chuyển đổi TensorFlow Lite để có hướng dẫn chi tiết về các kiến ​​thức cơ bản về chuyển đổi mô hình.

# 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

Sự suy luận

Để trình thông dịch TensorFlow Lite đọc đúng mô hình của bạn có chứa các toán tử TensorFlow Text, bạn phải định cấu hình nó để sử dụng các toán tử tùy chỉnh này và cung cấp các phương thức đăng ký cho chúng. Sử dụng tf_text.tflite_registrar.SELECT_TFTEXT_OPS để cung cấp bộ chức năng đăng ký đầy đủ cho các toán tử TensorFlow Text được hỗ trợ cho InterpreterWithCustomOps .

Lưu ý rằng trong khi ví dụ bên dưới hiển thị suy luận bằng Python, các bước tương tự trong các ngôn ngữ khác với một số bản dịch API nhỏ và sự cần thiết phải xây dựng tflite_registrar thành tệp nhị phân của bạn. Xem TensorFlow Lite Inference để biết thêm chi tiết.

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

Tiếp theo, trình thông dịch TensorFlow Lite được gọi với đầu vào, cung cấp một kết quả khớp với kết quả TensorFlow từ phía trên.

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