Cloud TPU v5e 推断转换器简介 [公开预览版]

预览

本产品或功能须遵守服务专用条款的“通用服务条款”部分中的“非正式版产品条款”。 正式发布前产品和功能“按原样”提供，可能仅提供有限支持。 如需了解详情，请参阅发布阶段说明。

简介

Cloud TPU Inference Converter 可以准备和优化 TensorFlow 2 (TF2) 模型以进行 TPU 推断。转换器在本地或 TPU 虚拟机 shell 中运行。建议使用 TPU 虚拟机 shell，因为它预安装了转换器所需的命令行工具。该工具会获取导出的 SavedModel 并执行以下步骤：

1. TPU 转换：它会向模型添加 TPUPartitionedCall 和其他 TPU 操作，使其可在 TPU 上运行。默认情况下，导出用于推断的模型没有此类操作，因此无法在 TPU 上服务，即使已在 TPU 上训练也是如此。
2. 批处理：它会向模型添加批处理操作，以实现图内批处理，以提高吞吐量。
3. BFloat16 转换：它会将模型的数据格式从 float32 转换为 bfloat16，以获得更好的计算性能并减少 TPU 上的高带宽内存 (HBM) 使用量。
4. IO 形状优化：它会针对在 CPU 和 TPU 之间传输的数据优化张量形状，以提高带宽利用率。

导出模型时，用户可以为要在 TPU 上运行的任何函数创建函数别名。它们将这些函数传递给 Converter，然后 Converter 将它们放在 TPU 上并进行优化。如果这些函数包含任何 TPU 不兼容的操作，则可以启用软设备放置，以便在 CPU 上执行不兼容的操作。

Cloud TPU Inference Converter 可以 Docker 映像提供，该映像可以在安装了 Docker 的任何环境中执行。

完成上述步骤的预计用时：约 20 分钟 - 30 分钟

前提条件

1. 模型必须是 TF2 模型，并以 SavedModel 格式导出。
2. 模型必须具有 TPU 函数的函数别名。如需了解如何执行此操作，请参阅代码示例。以下示例使用 tpu_func 作为 TPU 函数别名。
3. 确保您的机器的 CPU 支持高级矢量扩展 (AVX) 指令，因为 TensorFlow 库（Cloud TPU 推断转换器的依赖项）已编译为使用 AVX 指令。大多数 CPU 都支持 AVX。
   1. 您可以运行 lscpu | grep avx 以检查是否支持 AVX 指令集。

须知事项

在开始设置之前，请执行以下操作：

• 创建新项目：在 Google Cloud 控制台的项目选择器页面上，选择或创建 Cloud 项目。

• 设置 Cloud TPU 虚拟机：使用 Google Cloud 控制台或 gcloud 创建新的 Cloud TPU 虚拟机，或使用现有的 Cloud TPU 虚拟机在 Cloud TPU 虚拟机上使用转换后的模型进行推断。

  • 确保 Cloud TPU 虚拟机映像基于 Tensorflow。例如 --version=tpu-vm-tf-2.11.0。
  • 转换后的模型将在此 Cloud TPU 虚拟机上加载并提供。

• 确保您拥有使用 Cloud TPU 推断转换器所需的命令行工具。您可以在本地安装 Google Cloud SDK 和 Docker，也可以使用默认安装了此软件的 Cloud TPU 虚拟机。您可以使用这些工具与 Converter 映像进行交互。

  您可以使用以下命令，通过 SSH 连接到 Cloud TPU 虚拟机：

  gcloud compute tpus tpu-vm ssh ${tpu-name} --zone ${zone} --project ${project-id}
  

环境设置

通过 TPU 虚拟机 shell 或本地 shell 设置您的环境。

TPU 虚拟机 Shell

• 在您的 TPU 虚拟机 shell 中，运行以下命令以允许非根 Docker 使用：

  sudo usermod -a -G docker ${USER}
  newgrp docker
  

• 初始化您的 Docker 凭据帮助程序：

  gcloud auth configure-docker \
    us-docker.pkg.dev
  

本地 Shell

在本地 shell 中，按照以下步骤设置环境：

• 安装 Cloud SDK，其中包含 gcloud 命令行工具。

• 安装 Docker：

• 允许非根用户使用 Docker：

  sudo usermod -a -G docker ${USER}
  newgrp docker
  

• 登录您的环境：

  gcloud auth login
  

• 初始化您的 Docker 凭据帮助程序：

  gcloud auth configure-docker \
      us-docker.pkg.dev
  

• 拉取 Inference Converter Docker 映像：

    CONVERTER_IMAGE=us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:2.13.0
    docker pull ${CONVERTER_IMAGE}
    

转换器图片

该图片用于执行一次性模型转换。请设置模型路径并根据您的需求调整转换器选项。用法示例部分提供了几种常见用例。

注意 ：在下面的代码中，无需在 docker 命令中指定 mount，即可将 Cloud Storage 存储桶用于 input_model_dir 和 output_model_dir。

docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      function_alias: "tpu_func"
    }
    batch_options {
      num_batch_threads: 2
      max_batch_size: 8
      batch_timeout_micros: 5000
      allowed_batch_sizes: 2
      allowed_batch_sizes: 4
      allowed_batch_sizes: 8
      max_enqueued_batches: 10
    }
'

在 Cloud TPU 虚拟机中使用转换后的模型进行推断

# Initialize the TPU
resolver = tf.distribute.cluster_resolver.TPUClusterResolver("local")
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)

# Load the model
model = tf.saved_model.load(${CONVERTED_MODEL_PATH})

# Find the signature function for serving
serving_signature = 'serving_default' # Change the serving signature if needed
serving_fn = model.signatures[serving_signature]
# Run the inference using requests.
results = serving_fn(**inputs)
logging.info("Serving results: %s", str(results))

用法示例

为 TPU 函数添加函数别名

1. 在模型中查找或创建一个函数，该函数封装您要在 TPU 上运行的所有内容。如果 @tf.function 不存在，请添加它。
2. 保存模型时，请提供如下所示的 SaveOptions，为 model.tpu_func 指定别名 func_on_tpu。
3. 您可以将此函数别名传递给转换器。

class ToyModel(tf.keras.Model):
  @tf.function(
      input_signature=[tf.TensorSpec(shape=[None, 10], dtype=tf.float32)])
  def tpu_func(self, x):
    return x * 1.0

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'func_on_tpu': model.tpu_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

转换具有多个 TPU 函数的模型

您可以在 TPU 上放置多个函数。只需创建多个函数别名并将其传入 converter_options_string 传递给转换器即可。

tpu_functions {
  function_alias: "tpu_func_1"
}
tpu_functions {
  function_alias: "tpu_func_2"
}

量化

重要提示 ：此功能目前处于实验阶段，可能会发生更改，恕不另行通知。在将量化模型用于生产环境之前，用户应仔细评估其性能。TensorFlow 中提供对量化的 TPU v5e 支持。

量化是一种技术，可降低用于表示模型参数的数字的精度。这样可以减小模型大小并加快计算速度。量化模型可提供推断吞吐量的提升，以及更小的内存使用量和存储空间大小，但代价是准确率会略有下降。

TensorFlow 中以 TPU 为目标的全新训练后量化功能基于 TensorFlow Lite 中用于定位移动设备和边缘设备的类似现有功能开发而成。如需从总体上详细了解量化，您可以查看 TensorFlow Lite 的文档。

量化概念

本部分定义了与使用推断转换器的量化相关的概念。

TPU 系统架构页面介绍了与其他 TPU 配置（例如切片、主机、芯片和 TensorCore）相关的概念。

• 训练后量化 (PTQ)：PTQ 是一种技术，可缩减神经网络模型的大小和计算复杂性，而不会显著影响其准确性。PTQ 的工作原理是将经过训练的模型的浮点权重和激活数转换为精度较低的整数（例如 8 位或 16 位整数）。这样可以显著减小模型大小并缩短推断延迟时间，但只会对准确率产生微小的损失。

• 校准：量化的校准步骤是收集神经网络模型的权重和激活函数获得的一系列值的统计信息的过程。此信息用于确定模型的量化参数，这些值将用于将浮点权重和激活转换为整数。

• 代表性数据集：用于量化的代表性数据集是一个小型数据集，表示模型的实际输入数据。该模型在量化的校准步骤中使用，以收集模型的权重和激活将采用的值范围的统计数据。代表性数据集应满足以下条件：

  • 它应该正确表示在推断期间模型的实际输入。这意味着，它应该涵盖模型在现实世界中可能看到的值范围。
  • 它应共同流经每个条件分支（如 tf.cond）分支（如果有）。这一点很重要，因为量化过程需要能够处理模型的所有可能输入，即使它们未在代表性数据集中明确表示也是如此。
  • 范围应足够大，以收集足够的统计信息并减少错误。一般来讲，建议使用 200 个以上的代表性样本。

  代表性数据集可以是训练数据集的子集，也可以是专门设计用于代表模型真实输入的单独数据集。具体选择使用哪个数据集取决于具体应用。

• 静态范围量化 (SRQ)：SRQ 会在校准步骤期间确定神经网络模型的权重和激活值一次的值范围。这意味着对模型的所有输入使用相同的值范围。这种方法可能不如动态范围量化，尤其是对于具有大量输入值的模型。 不过，与动态范围量化相比，静态范围量化在运行时需要的计算量更少。

• 动态范围量化 (DRQ)：DRQ 为每个输入确定神经网络模型的权重和激活值的范围。这样一来，模型就可以适应输入数据值的范围，从而提高准确性。但是，与静态范围量化相比，动态范围量化在运行时需要更多的计算。

  功能	静态范围量化	动态范围量化
  值范围	在校准期间确定一次	针对每项输入确定
  准确性	可能不太准确，尤其是对于具有多种输入值的模型	可以更准确，尤其是对于具有多种输入值的模型
  复杂性	更简单	更复杂的
  运行时计算	计算更少	更多计算

• 仅权重量化：仅权重量化是一种量化，它仅量化神经网络模型的权重，而激活以浮点数表示。对于对准确性敏感的模型而言，这是一个不错的选择，因为这有助于保持模型的准确性。

如何使用量化

可以通过配置 QuantizationOptions 并将其设置为转换器选项来应用量化。重要选项包括：

• tags：用于标识要量化的 SavedModel 中的 MetaGraphDef 的标记集合。如果您只有一个 MetaGraphDef，则无需指定。
• sign_keys：标识包含输入和输出的 SignatureDef 的键序列。如果未指定，则使用 ["serving_default"]。
• quantization_method：要应用的量化方法。如果未指定，将应用 STATIC_RANGE 量化。
• op_set：应保留为 XLA。目前是默认选项，无需指定。
• 代表数据集：指定用于校准量化参数的数据集。

构建代表性数据集

代表性数据集实质上是可迭代样本。其中，样本是 {input_key: input_value} 的映射。例如：

representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]

应使用 tf-nightly pip 软件包当前提供的 TfRecordRepresentativeDatasetSaver 类将代表性数据集保存为 TFRecord 文件。例如：

# Assumed tf-nightly installed.
import tensorflow as tf
representative_dataset = [{"x": tf.random.uniform(shape=(3, 3))}
                          for _ in range(256)]
tf.quantization.experimental.TfRecordRepresentativeDatasetSaver(
       path_map={'serving_default': '/tmp/representative_dataset_path'}
    ).save({'serving_default': representative_dataset})

示例

以下示例使用签名键 serving_default 和函数别名 tpu_func 来量化模型：

注意 ：在以下代码中，将 CONVERTER_IMAGE 设置为 us-docker.pkg.dev/cloud-tpu-images/inference/tpu-inference-converter-cli:nightly 以使用每夜版 build。

docker run \
  --mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
  --mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
  ${CONVERTER_IMAGE} \
  --input_model_dir=/tmp/input \
  --output_model_dir=/tmp/output \
  --converter_options_string=' \
    tpu_functions { \
      function_alias: "tpu_func" \
    } \
    external_feature_configs { \
      quantization_options { \
        signature_keys: "serving_default" \
        representative_datasets: { \
          key: "serving_default" \
          value: { \
            tfrecord_file_path: "${TF_RECORD_FILE}" \
          } \
        } \
      } \
    } '

启用软设备放置

软设备放置可以使 TPU 不兼容的操作在运行时在 CPU 上运行。如果模型的 TPU 函数内部具有无法轻松分离的 TPU 不兼容的操作，您可以考虑通过为 TPU 函数设置 enable_soft_device_placement=True 来启用软设备放置。此标志将指示运行时检查是否存在 TPU 不兼容的操作并在 CPU 上执行。

注意 ：由于会影响性能，不建议将软设备放置用于推断用途。只应视为万不得已的方法。

软设备放置位置也有助于调试。因此，您可以在 TPU 函数中使用与 TPU 不兼容的 tf.print 来输出张量的值。

软设备放置与 MLIR 桥搭配使用时效果最佳。 需要设置一个额外的标志。

tpu_functions {
  function_alias: "tpu_func"
  enable_soft_device_placement: true
}

添加批处理

Converter 可用于向模型添加批处理。如需了解可调整的批处理选项的说明，请参阅批处理选项的定义。

默认情况下，Converter 将批量处理模型中的所有 TPU 函数。此外，它还可以批量处理用户提供的签名和函数，从而进一步提高性能。任何进行批处理的 TPU 函数、用户提供的函数或签名都必须满足批处理操作的严格的形状要求。

Converter 还可以更新现有的批处理选项。 以下示例说明了如何向模型添加批处理。如需详细了解批处理，请参阅批处理深入探究。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

停用 bfloat16 和 IO 形状优化

默认情况下，BFloat16 和 IO 形状优化处于启用状态。如果它们与您的模型不兼容，则可以将其停用。

# Disable both optimizations
disable_default_optimizations: true

# Or disable them individually
io_shape_optimization: DISABLED
bfloat16_optimization: DISABLED

转化报告

您可以在运行推断转换器后从日志中找到此转换报告。下面提供了一个示例。

-------- Conversion Report --------
TPU cost of the model: 96.67% (2034/2104)
CPU cost of the model:  3.33% (70/2104)

Cost breakdown
================================
%         Cost    Name
--------------------------------
3.33      70      [CPU cost]
48.34     1017    tpu_func_1
48.34     1017    tpu_func_2
--------------------------------

此报告会估算输出模型在 CPU 和 TPU 上的计算费用，并将 TPU 费用进一步细分为每个函数，从而反映您在转换器选项中选择的 TPU 函数。

如果您想更好地利用 TPU，可能需要试验模型结构并调整转换器选项。

注意 ：在当前转换报告中，TPU 函数中的所有操作均计为 TPU 费用。因此，如果模型具有软设备放置，则这些外部编译操作的费用计为 TPU 费用，而不是 CPU 费用。

常见问题解答

我应该在 TPU 上放置哪个/哪些函数？

最好将尽可能多的模型放在 TPU 上，因为绝大多数操作在 TPU 上执行得更快。

如果您的模型不包含任何与 TPU 不兼容的操作、字符串或稀疏张量，那么将整个模型放在 TPU 上通常是最佳策略。为此，您可以查找或创建一个封装整个模型的函数，为其创建函数别名，然后将其传递给 Converter。

如果模型包含无法在 TPU 上工作的部分（例如，TPU 不兼容的操作、字符串或稀疏张量），则 TPU 函数的选择取决于不兼容部分的位置。

• 如果它位于模型的开头或结尾，则可以重构模型以将其保留在 CPU 上。例如，字符串预处理和后处理阶段。如需详细了解如何将代码移至 CPU，请参阅“如何将模型的一部分移至 CPU？”它显示了重构模型的典型方法。
• 如果它在模型中间，并且只有几次操作，则可以考虑使用软设备放置。但是，由于 CPU 与 TPU 之间需要额外的通信，它可能会产生较长的推断延迟时间。
• 如果它位于模型中间，并且是一个很大的运算，那么软设备放置会使推断成本过高。因此，最好将模型拆分为三部分，并在中间包含所有 TPU 不兼容的操作，并使其在 CPU 上运行。
• 如果它是一个稀疏张量，请考虑对 CPU 调用 tf.sparse.to_dense，并将生成的密集张量传递给模型的 TPU 部分。

另一个需要考虑的因素是 HBM 使用情况。嵌入表可能会使用大量 HBM。如果它们超出 TPU 的硬件限制，则必须与查找操作一起放在 CPU 上。

如有可能，一个签名下只能有一个 TPU 函数。如果模型的结构需要为每个传入推断请求调用多个 TPU 函数，您应该了解 CPU 和 TPU 之间发送张量的延迟时间会增加。

要评估 TPU 函数的选择情况，查看转换报告是一个好方法。它显示了放置在 TPU 上的计算百分比，以及每个 TPU 函数的成本明细。

如何将模型的一部分移至 CPU？

如果模型包含无法在 TPU 上处理的部分，您需要重构模型才能将这些部分迁移到 CPU。这是一个玩具示例。该模型是具有预处理阶段的语言模型。为简单起见，我们省略了层定义和函数的代码。

class LanguageModel(tf.keras.Model):
  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.bert_layer(word_ids)

此模型无法直接在 TPU 上提供，原因有两个。首先，参数是一个字符串。其次，preprocess 函数可能包含许多字符串操作。两者都不兼容 TPU。

如需重构此模型，您可以再创建一个名为 tpu_func 的函数来托管计算密集型 bert_layer。然后，为 tpu_func 创建函数别名，并将其传递给 Converter。这样，tpu_func 中的所有内容都将在 TPU 上运行，而 model_func 中剩余的所有内容都将在 CPU 上运行。

class LanguageModel(tf.keras.Model):
  @tf.function
  def tpu_func(self, word_ids):
    return self.bert_layer(word_ids)

  @tf.function
  def model_func(self, input_string):
    word_ids = self.preprocess(input_string)
    return self.tpu_func(word_ids)

如果模型具有与 TPU 不兼容的操作、字符串或稀疏张量，该怎么办？

TPU 支持大多数标准 TensorFlow 操作，但也不支持一些（包括稀疏张量和字符串）操作。转换器不会检查 TPU 不兼容的操作。因此，包含此类操作的模型可以传递转化。但在运行推断模型时，将会发生如下错误。

'tf.StringToNumber' op isn't compilable for TPU device. enable soft_device_placement option to run on CPU

如果模型具有与 TPU 不兼容的操作，则应将它们放在 TPU 函数之外。此外，字符串是 TPU 上不受支持的数据格式。因此，不应将字符串类型的变量放置在 TPU 函数中。TPU 函数的参数和返回值也不应是字符串类型。同样，请避免将稀疏张量放入 TPU 函数中，包括其参数和返回值中。

重构模型的不兼容部分并将其迁移到 CPU 通常并不难。请参考这个示例。 但在某些情况下，与 TPU 不兼容的操作会深层嵌套在模型中，并且难以分离。对于这些型号，您可以考虑启用软设备放置，以在运行时在 CPU 上运行操作。但请注意性能影响。

如何在模型中支持自定义操作？

如果在模型中使用自定义操作，转换器可能无法识别这些操作，并且无法转换模型。这是因为自定义操作的操作库（包含操作的完整定义）未关联到 Converter。

由于转换器代码目前尚未开源，因此无法使用自定义操作构建转换器。

如果我有 TensorFlow 1 模型，该怎么办？

Converter 不支持 TensorFlow 1 模型。TensorFlow 1 模型应迁移到 TensorFlow 2。

运行模型时，我需要启用 MLIR 桥吗？

大多数转换后的模型都可以使用较新的 TF2XLA MLIR 桥或原始 TF2XLA 桥运行。

在训练和推断期间，TensorFlow 2 中的软设备放置与 MLIR 桥搭配使用效果最佳。可以使用这些说明启用 MLIR 桥。

如何转换导出时没有函数别名的模型？

如果导出模型时没有指定函数别名，最简单的方法是再次导出该模型并创建函数别名。如果无法重新导出，仍然可以通过提供 concrete_function_name 来转换模型。不过，识别正确的 concrete_function_name 需要进行一些检测工作。

函数别名是从用户定义的字符串到具体函数名称的映射。它们可让您更轻松地引用模型中的特定函数。转换器既接受函数别名，也接受原始的具体函数名称。

您可以通过检查 saved_model.pb 找到具体的函数名称。

以下示例展示了如何将一个名为 __inference_serve_24 的具体函数放置在 TPU 上。

sudo docker run \
--mount type=bind,source=${MODEL_PATH},target=/tmp/input,readonly \
--mount type=bind,source=${CONVERTED_MODEL_PATH},target=/tmp/output \
${CONVERTER_IMAGE} \
--input_model_dir=/tmp/input \
--output_model_dir=/tmp/output \
--converter_options_string='
    tpu_functions {
      concrete_function_name: "__inference_serve_24"
    }'

如何解决编译时间常量约束错误？

对于训练和推断，XLA 要求某些操作的输入在 TPU 编译时具有已知形状。这意味着，当 XLA 编译程序的 TPU 部分时，这些操作的输入必须具有静态已知的形状。

您可以通过两种方式解决此问题。

• 最佳方式是在 XLA 编译 TPU 程序时更新操作的输入，使其具有静态已知的形状。此编译会在运行模型的 TPU 部分之前立即发生。这意味着，当 TpuFunction 即将运行时，形状应该是静态已知的。
• 另一种方法是修改 TpuFunction，使其不再包含有问题的操作。

为什么我会收到批处理形状错误？

批处理具有严格的形状要求，允许按其第 0 个维度（也称为批处理维度）对传入请求进行批处理。这些形状要求来自 TensorFlow 批处理操作，无法放宽。

不满足这些要求将导致错误，例如：

1. 批处理输入张量必须至少有一个维度。
2. 输入维度应匹配。
3. 在给定操作调用中提供的批处理输入张量必须具有相同的 0 维大小。
4. 批量输出张量的第 0 个维度不等于输入张量的第 0 个维度大小的总和。

如需满足这些要求，请考虑提供其他要批处理的函数或签名。您可能还需要修改现有函数以满足这些要求。

如果对一个函数执行批处理，请确保其 @tf.function 的 input_signature 的形状在第 0 个维度中均具有 None。如果要对签名进行批处理，请确保其所有输入在第 0 个维度中均为 -1。

如需全面了解发生这些错误的原因以及如何解决这些问题，请参阅深入探究批处理。

已知问题

TPU 函数无法间接调用其他 TPU 函数

虽然转换器可以处理跨 CPU-TPU 边界的大多数函数调用场景，但在一种极少数情况下会失败。这种情况是指 TPU 函数间接调用另一个 TPU 函数。

这是因为转换器会将 TPU 函数的直接调用方从调用 TPU 函数本身修改为调用 TPU 调用存根。调用桩包含只能在 CPU 上工作的操作。当 TPU 函数调用任何最终调用直接调用方的函数时，这些 CPU 操作可以在 TPU 上执行，从而产生缺少内核错误。请注意，此情况与直接调用另一个 TPU 函数的 TPU 函数不同。在这种情况下，转换器不会修改任一函数来调用调用桩，因此它可以正常运行。

在 Converter 中，我们实现了针对此场景的检测。如果您看到以下错误，则表示您的模型已遇到这种极端情况：

Unable to place both "__inference_tpu_func_2_46" and "__inference_tpu_func_4_68"
on the TPU because "__inference_tpu_func_2_46" indirectly calls
"__inference_tpu_func_4_68". This behavior is unsupported because it can cause
invalid graphs to be generated.

一般解决方案是重构模型，以避免此类函数调用场景。如果您觉得这很难实现，请与 Google 支持团队联系，以了解详情。

参考文档

Protobuf 格式的转换器选项

message ConverterOptions {
  // TPU conversion options.
  repeated TpuFunction tpu_functions = 1;

  // The state of an optimization.
  enum State {
    // When state is set to default, the optimization will perform its
    // default behavior. For some optimizations this is disabled and for others
    // it is enabled. To check a specific optimization, read the optimization's
    // description.
    DEFAULT = 0;
    // Enabled.
    ENABLED = 1;
    // Disabled.
    DISABLED = 2;
  }

  // Batch options to apply to the TPU Subgraph.
  //
  // At the moment, only one batch option is supported. This field will be
  // expanded to support batching on a per function and/or per signature basis.
  //
  //
  // If not specified, no batching will be done.
  repeated BatchOptions batch_options = 100;

  // Global flag to disable all optimizations that are enabled by default.
  // When enabled, all optimizations that run by default are disabled. If a
  // default optimization is explicitly enabled, this flag will have no affect
  // on that optimization.
  //
  // This flag defaults to false.
  bool disable_default_optimizations = 202;

  // If enabled, apply an optimization that reshapes the tensors going into
  // and out of the TPU. This reshape operation improves performance by reducing
  // the transfer time to and from the TPU.
  //
  // This optimization is incompatible with input_shape_opt which is disabled.
  // by default. If input_shape_opt is enabled, this option should be
  // disabled.
  //
  // This optimization defaults to enabled.
  State io_shape_optimization = 200;

  // If enabled, apply an optimization that updates float variables and float
  // ops on the TPU to bfloat16. This optimization improves performance and
  // throughtput by reducing HBM usage and taking advantage of TPU support for
  // bfloat16.
  //
  // This optimization may cause a loss of accuracy for some models. If an
  // unacceptable loss of accuracy is detected, disable this optimization.
  //
  // This optimization defaults to enabled.
  State bfloat16_optimization = 201;

  BFloat16OptimizationOptions bfloat16_optimization_options = 203;

  // The settings for XLA sharding. If set, XLA sharding is enabled.
  XlaShardingOptions xla_sharding_options = 204;
}

message TpuFunction {
  // The function(s) that should be placed on the TPU. Only provide a given
  // function once. Duplicates will result in errors. For example, if
  // you provide a specific function using function_alias do not also provide the
  // same function via concrete_function_name or jit_compile_functions.
  oneof name {
    // The name of the function alias associated with the function that
    // should be placed on the TPU. Function aliases are created during model
    // export using the tf.saved_model.SaveOptions.
    //
    // This is a recommended way to specify which function should be placed
    // on the TPU.
    string function_alias = 1;

    // The name of the concrete function that should be placed on the TPU. This
    // is the name of the function as it found in the GraphDef and the
    // FunctionDefLibrary.
    //
    // This is NOT the recommended way to specify which function should be
    // placed on the TPU because concrete function names change every time a
    // model is exported.
    string concrete_function_name = 3;

    // The name of the signature to be placed on the TPU. The user must make
    // sure there is no TPU-incompatible op under the entire signature. Or soft
    // device placement can be enabled. But this will incur performance loss.
    string signature_name = 5;

    // When jit_compile_functions is set to True, all jit compiled functions
    // are placed on the TPU.
    //
    // To use this option, decorate the relevant function(s) with
    // @tf.function(jit_compile=True), before exporting. Then set this flag to
    // True. The converter will find all functions that were tagged with
    // jit_compile=True and place them on the TPU.
    //
    // When using this option, all other settings for the TpuFunction such as
    // enable_soft_device_placement will apply to all functions tagged with
    // jit_compile=True.
    //
    // This option will place all jit_compile=True functions on the TPU.
    // If only some jit_compile=True functions should be placed on the TPU,
    // use function_alias or concrete_function_name.
    bool jit_compile_functions = 4;
  }

  // Set to true to enable outside compilation for this TPU function. If the TPU
  // function has TPU-incompatible ops, outside compilation can automatically
  // move the ops to execute on the CPU at the runtime.
  bool enable_soft_device_placement = 2;
}

message BatchOptions {
  // Number of scheduling threads for processing batches of work. Determines
  // the number of batches processed in parallel. This should be roughly in line
  // with the number of TPU cores available.
  int32 num_batch_threads = 1;

  // The maximum allowed batch size.
  int32 max_batch_size = 2;

  // Maximum number of microseconds to wait before outputting an incomplete
  // batch.
  int32 batch_timeout_micros = 3;

  // Optional list of allowed batch sizes. If left empty,
  // does nothing. Otherwise, supplies a list of batch sizes, causing the op
  // to pad batches up to one of those sizes. The entries must increase
  // monotonically, and the final entry must equal max_batch_size.
  repeated int32 allowed_batch_sizes = 4;

  // Maximum number of batches enqueued for processing before requests are
  // failed fast.
  int32 max_enqueued_batches = 5;

  // If set, disables large batch splitting which is an efficiency improvement
  // on batching to reduce padding inefficiency.
  bool disable_large_batch_splitting = 6;

  // Experimental features of batching. Everything inside is subject to change.
  message Experimental {
    // The component to be batched.
    // 1. Unset if it's for all TPU subgraphs.
    // 2. Set function_alias or concrete_function_name if it's for a function.
    // 3. Set signature_name if it's for a signature.
    oneof batch_component {
      // The function alias associated with the function. Function alias is
      // created during model export using the tf.saved_model.SaveOptions, and is
      // the recommended way to specify functions.
      string function_alias = 1;

      // The concreate name of the function. This is the name of the function as
      // it found in the GraphDef and the FunctionDefLibrary. This is NOT the
      // recommended way to specify functions, because concrete function names
      // change every time a model is exported.
      string concrete_function_name = 2;

      // The name of the signature.
      string signature_name = 3;
    }
  }

  Experimental experimental = 7;
}

message BFloat16OptimizationOptions {
  // Indicates where the BFloat16 optimization should be applied.
  enum Scope {
    // The scope currently defaults to TPU.
    DEFAULT = 0;
    // Apply the bfloat16 optimization to TPU computation.
    TPU = 1;
    // Apply the bfloat16 optimization to the entire model including CPU
    // computations.
    ALL = 2;
  }

  // This field indicates where the bfloat16 optimization should be applied.
  //
  // The scope defaults to TPU.
  Scope scope = 1;

  // If set, the normal safety checks are skipped. For example, if the model
  // already contains bfloat16 ops, the bfloat16 optimization will error because
  // pre-existing bfloat16 ops can cause issues with the optimization. By
  // setting this flag, the bfloat16 optimization will skip the check.
  //
  // This is an advanced feature and not recommended for almost all models.
  //
  // This flag is off by default.
  bool skip_safety_checks = 2;

  // Ops that should not be converted to bfloat16.
  // Inputs into these ops will be cast to float32, and outputs from these ops
  // will be cast back to bfloat16.
  repeated string filterlist = 3;
}

message XlaShardingOptions {
  // num_cores_per_replica for TPUReplicateMetadata.
  //
  // This is the number of cores you wish to split your model into using XLA
  // SPMD.
  int32 num_cores_per_replica = 1;

  // (optional) device_assignment for TPUReplicateMetadata.
  //
  // This is in a flattened [x, y, z, core] format (for
  // example, core 1 of the chip
  // located in 2,3,0 will be stored as [2,3,0,1]).
  //
  // If this is not specified, then the device assignments will utilize the same
  // topology as specified in the topology attribute.
  repeated int32 device_assignment = 2;

  // A serialized string of tensorflow.tpu.TopologyProto objects, used for
  // the topology attribute in TPUReplicateMetadata.
  //
  // You must specify the mesh_shape and device_coordinates attributes in
  // the topology object.
  //
  // This option is required for num_cores_per_replica > 1 cases due to
  // ambiguity of num_cores_per_replica, for example,
  // pf_1x2x1 with megacore and df_1x1
  // both have num_cores_per_replica = 2, but topology is (1,2,1,1) for pf and
  // (1,1,1,2) for df.
  // - For pf_1x2x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,2,1,1]
  //   device_coordinates=flatten([0,0,0,0], [0,1,0,0])
  // - For df_1x1, mesh shape and device_coordinates looks like:
  //   mesh_shape = [1,1,1,2]
  //   device_coordinates=flatten([0,0,0,0], [0,0,0,1])
  // - For df_2x2, mesh shape and device_coordinates looks like:
  //   mesh_shape = [2,2,1,2]
  //   device_coordinates=flatten(
  //    [0,0,0,0],[0,0,0,1],[0,1,0,0],[0,1,0,1]
  //    [1,0,0,0],[1,0,0,1],[1,1,0,0],[1,1,0,1])
  bytes topology = 3;
}

批处理深入介绍

批处理用于提高吞吐量和 TPU 利用率。它允许同时处理多个请求。在训练期间，可以使用 tf.data 完成批处理。在推断期间，通常是通过在图中添加对传入请求进行批处理的操作来完成的。该操作会等到有足够的请求或达到超时之后再根据单个请求生成大批量。如需详细了解可以调整的不同批处理选项（包括批次大小和超时），请参阅批处理选项的定义。

默认情况下，转换器会直接在 TPU 计算之前插入批处理操作。它会使用批处理操作封装用户提供的 TPU 函数和模型中原有的所有 TPU 计算。您可以替换此默认行为，方法是告知 Converter 应批量处理哪些函数和/或签名。

以下示例展示了如何添加默认批处理。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}

签名批量处理

注意 ：此功能仍处于实验阶段。测试完成后，该字段将从实验性封装容器中移出。

签名批处理从签名的输入开始批量处理整个模型，然后一直到签名的输出。与 Converter 的默认批处理行为不同，签名批处理会同时对 TPU 计算和 CPU 计算进行批处理。这样一来，在某些模型上进行推断期间，性能会提升 10% 到 20%。

与所有批处理一样，签名批处理具有严格的形状要求。为了确保满足这些形状要求，签名输入应具有至少具有两个维度的形状。第一个维度是批次大小，大小应为 -1。例如，(-1, 4)、(-1) 或 (-1, 128, 4, 10) 都是有效的输入形状。如果无法做到这一点，请考虑使用默认批处理行为或函数批处理。

如需使用签名批量处理功能，请使用 BatchOptions 以 signature_name 的形式提供签名名称。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    signature_name: "serving_default"
  }
}

函数批处理

注意 ：此功能仍处于实验阶段。测试完成后，该字段将从实验性封装容器中移出。

函数批量处理可用于告知 Converter 应批量处理哪些函数。默认情况下，Converter 会批处理所有 TPU 函数。函数批处理功能会覆盖此默认行为。

函数批处理可用于批量处理 CPU 计算。对许多模型进行批量 CPU 计算时，性能都会得到提升。批处理 CPU 计算的最佳方式是使用签名批处理，但它可能不适用于某些模型。在这些情况下，函数批处理除了 TPU 计算之外，还可用于批处理部分 CPU 计算。请注意，批处理操作无法在 TPU 上运行，因此必须在 CPU 上调用所提供的任何批处理函数。

函数批处理还可用于满足批处理操作规定的严格的形状要求。如果 TPU 函数不符合批处理操作的形状要求，则可以使用函数批处理来告知转换器对不同的函数进行批处理。

为此，请为应进行批处理的函数生成 function_alias。为此，您可以在模型中查找或创建一个函数，该函数封装所有要批处理的所有内容。确保此函数符合批处理操作严格的形状要求要求。如果还没有，请添加 @tf.function。请务必向 @tf.function 提供 input_signature。第 0 个维度应为 None，因为它是批次维度，因此它不能是固定大小。例如，[None, 4]、[None] 或 [None, 128, 4, 10] 都是有效的输入形状。保存模型时，请提供如下所示的 SaveOptions，为 model.batch_func 指定别名“batch_func”。然后，您可以将此函数别名传递给转换器。

class ToyModel(tf.keras.Model):
  @tf.function(input_signature=[tf.TensorSpec(shape=[None, 10],
                                              dtype=tf.float32)])
  def batch_func(self, x):
    return x * 1.0

  ...

model = ToyModel()
save_options = tf.saved_model.SaveOptions(function_aliases={
    'batch_func': model.batch_func,
})
tf.saved_model.save(model, model_dir, options=save_options)

接下来，使用 BatchOptions 传递 function_alias。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
  experimental {
    function_alias: "batch_func"
  }
}

批处理选项的定义

• num_batch_threads：（整数）用于处理工作批次的调度线程数。确定并行处理的批次数量。这应该与可用的 TPU 核心数量大致一致。
• max_batch_size：（整数）允许的批次大小上限。该值可以大于 allowed_batch_sizes，以便利用大批量拆分。
• batch_timeout_micros：（整数）在输出未完成的批次之前等待的最大微秒数。
• allowed_batch_sizes：（整数列表）如果列表不为空，则会填充批次，直到列表中最接近的大小。该列表必须单调递增，并且最终元素必须小于或等于 max_batch_size。
• max_enqueued_batches：（整数）在请求快速失败之前加入队列等待处理的批次数上限。

更新现有批处理选项

如需添加或更新批处理选项，您可以添加或更新批处理选项，方法是运行指定 batch_options 的 Docker 映像，并使用 --converter_options_string 标志将 disable_default_optimizations 设置为 true。批处理选项将应用于每个 TPU 函数或现有的批处理操作。

注意 ：您无法通过同时使用 --converter_options_string 设置 tpu_functions 来更新批处理选项。

batch_options {
  num_batch_threads: 2
  max_batch_size: 8
  batch_timeout_micros: 5000
  allowed_batch_sizes: 2
  allowed_batch_sizes: 4
  allowed_batch_sizes: 8
  max_enqueued_batches: 10
}
disable_default_optimizations=True

批处理形状要求

批次的创建方式是沿着其批次（第 0 个）维度串联各个请求的输入张量。输出张量沿其第 0 个维度拆分。为了执行这些操作，批处理操作对其输入和输出有严格的形状要求。

演示

如需了解这些要求，最好先了解如何执行批处理。在下面的示例中，我们将批处理一个简单的 tf.matmul 操作。

def my_func(A, B)
    return tf.matmul(A, B)

第一个推断请求生成形状分别为 (1, 3, 2) 和 (1, 2, 4) 的输入 A 和 B。第二个推断请求生成形状 (2, 3, 2) 和 (2, 2, 4) 的输入 A 和 B。

已达到批处理超时时间。该模型支持 3 的批次大小，因此推断请求 1 和 2 会一起进行批处理，而不会进行任何填充。批处理张量通过沿批处理（第 0 个）维度串联请求 1 和 2 来形成。由于 #1 的 A 的形状为 (1, 3, 2)，#2 的 A 的形状为 (2, 3, 2)，因此当它们沿批次（第 0）维度串联时，生成的形状为 (3, 3, 2)。

执行 tf.matmul 后，它会生成形状为 (3, 3, 4) 的输出。

tf.matmul 的输出是批量处理的，因此需要将其拆分回单独的请求。批处理操作通过沿每个输出张量的批处理（第 0 个）维度进行拆分来实现此目的。它决定如何根据原始输入的形状拆分第 0 个维度。由于请求 1 的形状的第 0 个维度为 1，因此对于形状为 (1, 3, 4)，其输出的第 0 个维度为 1。由于请求 2 的形状的第 0 个维度为 2，因此对于形状为 (2, 3, 4)，其输出的第 0 个维度为 2。

形状要求

为了执行上述输入串联和输出拆分，批处理操作具有以下形状要求：

1. 批处理的输入不能是标量。要沿第 0 个维度串联，张量必须至少有两个维度。

   上述演示中。A 和 B 都不是标量。

   不满足此要求将导致错误，例如：Batching input tensors must have at least one dimension。解决此问题的一个简单方法是将标量设为向量。

2. 在不同的推断请求（例如，不同的会话运行调用）中，除了第 0 个维度之外，同名的输入张量在每个维度上都具有相同的大小。这样，输入可以沿其第 0 个维度正常串联。

   在上面的演示中，请求 1 的 A 的形状为 (1, 3, 2)。这意味着，将来的任何请求都必须生成格式为 (X, 3, 2) 的形状。请求 2 通过 (2, 3, 2) 满足此要求。同样，请求 1 的 B 的形状为 (1, 2, 4)，因此未来的所有请求都必须生成一个具有 (X, 2, 4) 模式的形状。

   不满足此要求将导致错误，例如：Dimensions of inputs should match。

3. 对于给定的推断请求，所有输入都必须具有相同的第 0 个维度大小。如果批处理操作的不同输入张量具有不同的第 0 个维度，则批处理操作不知道如何拆分输出张量。

   在上面的演示中，请求 1 的张量都具有第 1 个维度大小为 1。这样，批处理操作就能知道其输出的第 0 个维度大小应为 1。类似地，请求 2 的张量的第 0 个维度大小为 2，因此其输出的第 0 个维度大小为 2。当批处理操作拆分 (3, 3, 4) 的最终形状时，它会针对请求 1 生成 (1, 3, 4)，针对请求 2 生成 (2, 3, 4)。

   不满足此要求将导致错误，例如：Batching input tensors supplied in a given op invocation must have equal 0th-dimension size。

4. 每个输出张量形状的第 0 个维度大小必须是所有输入张量的第 0 个维度大小的总和（加上批处理操作引入的任何填充，以满足下一个最大的 allowed_batch_size）。这样，批处理操作就可以根据输入张量的第 0 个维度沿其第 0 个维度拆分输出张量。

   在上面的演示中，输入张量的第 0 个维度为 1（来自请求 1）和 2（来自请求 2）。因此，每个输出张量的第 0 个维度必须为 3，因为 1+2=3。输出张量 (3, 3, 4) 满足此要求。如果 3 不是有效的批次大小，但 4 是有效的，则批处理操作必须填充输入的第 0 维度（从 3 到 4）。在这种情况下，每个输出张量的第 0 个维度大小必须为 4。

   不满足此要求将导致错误，例如：Batched output tensor's 0th dimension does not equal the sum of the 0th dimension sizes of the input tensors。

解决形状要求错误

如需满足这些要求，请考虑提供其他要批处理的函数或签名。您可能还需要修改现有函数以满足这些要求。

如果要对某个函数进行批处理，请确保其 @tf.function 的 input_signature 形状在第 0 个维度（也称为批量维度）中都具有 None。如果要批量处理签名，请确保其所有输入在第 0 个维度中均为 -1。

BatchFunction 操作不支持将 SparseTensors 作为输入或输出。在内部，每个稀疏张量都表示为三个独立的张量，可以具有不同的第 0 个维度大小。