支持的精选 TensorFlow 算子

小心:算子列表会经常更新。

TensorFlow 核心算子

以下是具有精选 TensorFlow 算子功能的 TensorFlow Lite 运行时支持的 TensorFlow 核心算子的详尽列表。

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 2.5 版开始提供

如果您创建了自己的 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 算子,则有三个地方可以像此提交一样更新。

(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/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`
...

因为上面的列表按字母顺序排序,因此能确保名称处于正确的位置。