支援的 Select TensorFlow 運算子

TensorFlow 核心運算子

以下是 TensorFlow 核心運算的完整清單,TensorFlow Lite 執行階段透過 Select TensorFlow Ops 功能提供支援。

TensorFlow Text 和 SentencePiece 運算子

如果您使用 Python API 進行轉換並匯入這些程式庫,則支援下列 TensorFlow TextSentencePiece 運算子。

TF.Text 運算子

  • CaseFoldUTF8
  • ConstrainedSequence
  • MaxSpanningTree
  • NormalizeUTF8
  • NormalizeUTF8WithOffsetsMap
  • RegexSplitWithOffsets
  • RougeL
  • SentenceFragments
  • SentencepieceOp
  • SentencepieceTokenizeOp
  • SentencepieceTokenizeWithOffsetsOp
  • SentencepieceDetokenizeOp
  • SentencepieceVocabSizeOp
  • SplitMergeTokenizeWithOffsets
  • UnicodeScriptTokenizeWithOffsets
  • WhitespaceTokenizeWithOffsets
  • WordpieceTokenizeWithOffsets

SentencePiece 運算子

  • SentencepieceGetPieceSize
  • SentencepiecePieceToId
  • SentencepieceIdToPiece
  • SentencepieceEncodeDense
  • SentencepieceEncodeSparse
  • SentencepieceDecode

以下程式碼片段示範如何轉換包含上述運算子的模型

import tensorflow as tf
# These imports are required to load operators' definition.
import tensorflow_text as tf_text
import sentencepiece as spm

converter = tf.lite.TFLiteConverter.from_keras_model(your_model)
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
model_data = converter.convert()

在執行階段方面,也需要將 TensorFlow Text 或 SentencePiece 程式庫連結到最終應用程式或二進位檔。

使用者定義的運算子

如果您建立了自己的 TensorFlow 運算子,您也可以透過在 experimental_select_user_tf_ops 中列出必要運算子,將包含這些運算子的模型轉換為 TensorFlow Lite,如下所示

import tensorflow as tf

ops_module = tf.load_op_library('./your_ops_library.so')

converter = tf.lite.TFLiteConverter.from_saved_model(your_model)
converter.target_spec.supported_ops = [
  tf.lite.OpsSet.TFLITE_BUILTINS, tf.lite.OpsSet.SELECT_TF_OPS
]
converter.target_spec.experimental_select_user_tf_ops = [
    'your_op_name1',
    'your_op_name2'
]
model_data = converter.convert()

在執行階段方面,也需要將您的運算子程式庫連結到最終應用程式或二進位檔。

將 TensorFlow 核心運算子新增至允許清單。

如果您遇到 TensorFlow 核心運算子不在上述允許清單中的情況,您可以在此處回報功能要求,並提供未在允許清單中列出的 TensorFlow 核心運算子名稱。

您也可以從原始碼建立自己的提取要求。例如,如果您想在允許清單中新增 raw_ops.StringToNumber 運算子,則有三個地方需要更新,例如這個 commit

(1) 將運算子核心原始碼新增至 portable_extended_ops_group2 BUILD 規則。

filegroup(
    name = "portable_extended_ops_group2",
    srcs = [
        ...
+       "string_to_number_op.cc",

        ...
    ],
)

為了在 tensorflow/core/kernels 目錄下找到相關的運算子核心原始碼檔案,您可以搜尋原始碼位置,其中包含具有運算子名稱的下列核心宣告

REGISTER_KERNEL_BUILDER(Name("StringToNumber")                 \
                            .Device(DEVICE_CPU)                \
                            .TypeConstraint<type>("out_type"), \
                        StringToNumberOp<type>)

如果 tensorflow/core/kernels 目錄下有任何標頭檔是運算子核心原始碼中需要的,則您需要將標頭檔新增至 portable_extended_ops_headers BUILD 規則,如下所示

filegroup(
    name = "portable_extended_ops_headers",
    srcs = [
        ...
+       "string_util.h",

        ...
    ],
)

(2) 將運算子名稱新增至允許清單。

允許清單在 tensorflow/compiler/mlir/lite/delegates/flex/allowlisted_flex_ops.cc 中定義。需要列出 TensorFlow 核心運算子名稱,才能透過 Select TF 選項允許使用。

static const std::set<std::string>* allowlisted_flex_ops =
    new std::set<std::string>({
        ...
+       "StringToNumber",

        ...
    });

由於上述清單依字母順序排序,因此請務必將名稱放在正確的位置。

(3) 將運算子名稱新增至本指南頁面。

為了向其他開發人員顯示運算子包含項目,也應更新本指南頁面。此頁面位於 tensorflow/lite/g3doc/guide/op_select_allowlist.md

## TensorFlow core operators

The following is an exhaustive list of TensorFlow core operations that are
supported by TensorFlow Lite runtime with the Select TensorFlow Ops feature.

...
+*   `raw_ops.StringToNumber`
...

由於上述清單依字母順序排序,因此請務必將名稱放在正確的位置。