性能指南

Cloud TPU 可提供高性能并且费用较低。通过针对您的应用调整 Cloud TPU 配置参数,以及找出并解决任何限制性能的瓶颈,您可以进一步增强 Cloud TPU 性能。

本指南为您介绍如何执行以下操作来最大限度地提高 Cloud TPU 性能:

此外,以下资源中说明了如何实现以下目的:

模型处理性能

本部分介绍可能会降低模型性能的常见问题,以及如何解决这些问题。

  1. 数据预处理花费的时间过多

    TensorFlow-TPU 软件堆栈允许用户在将数据馈送到 TPU 之前,在 CPU 上执行复杂的数据预处理。由于 TPU 的运行速度非常快,因此如果输入数据非常大或以复杂的方式进行处理,则输入数据处理很容易成为瓶颈。

    Cloud TPU 剖析工具提供了一种简单的方法来衡量输入处理是否是模型的瓶颈。在这种情况下,请以一次性费用离线执行昂贵的预处理,而不是在它们训练的每个模型的每个周期产生费用。

  2. 由于分片,批次大小太小(跨核心拆分批次)

    如果您使用 n 的训练批次大小,则将跨 8 个 TPU 核心对此项进行拆分或分片。如果 n = 128 且您将 TPU 与单个 P100 进行比较,则 GPU 可能会以峰值 FLOPs 的 70% 运行,而 TPU 仅以峰值 FLOPs 的 20% 运行(因为分片 TPU 批次大小实际上为 16,(128 / 8)),而非 128。在此示例中,批次大小为 1024 的 TPU 会更好地利用 TPU 内存。

    为获得最佳内存用量,请使用内存中可容纳的最大批次大小。每个 TPU 核心都使用 128 x 128 内存单元矩阵进行处理。通常,您的批次大小应该能被 128 整除,才能有效地使用 TPU 内存。

XLA 编译器性能

XLA 是一款机器学习编译器,可以为 TPU、CPU、GPU 和其他平台生成二进制文件。它是标准 TensorFlow 代码库的一部分。 适用于 Cloud TPU 的 TensorFlow 模型被转换为 XLA 图,然后 XLA 将其编译为 TPU 可执行文件。 如需详细了解 XLA 和 TensorFlow 如何交互,请参阅 XLA 概览

Cloud TPU 硬件不同于 CPU 和 GPU。总体来说,CPU 的特征是具有少数高性能线程。而 GPU 则具有大量低性能线程。Cloud TPU 具有 128 x 128 矩阵单元,既可作为能在每个周期执行 16K 操作的单个强大线程,也可作为 128 x 128 个以流水线形式连接的小型简单线程。因此,在内存寻址过程中,最好为其使用 8 的倍数(浮点数),而针对矩阵单元的操作最好为 128 的倍数。

平铺的影响

Cloud TPU 中会平铺数组。这需要将其中一个维度填充为 8 的倍数,并将另一个维度填充为 128 的倍数。 XLA 将执行数据布局转换,并且将数据安排在内存中,以使硬件高效地处理数据。这一转换是通过启发法推进的。虽然编译器在大多数情况下具有出色表现,但也有可能出错。为了实现最高性能,尝试不同的模型配置可能会有所帮助。

经验法则:尽可能降低填充成本

使用平铺内存方案的一个后果是:能否高效地利用内存取决于在填充开销上浪费的内存量。 要以最高效的方式使用 Cloud TPU,请使用可最小化平铺开销的维度大小。

例如,卷积结果包括 (1) 批量、(2) 输出特征和 (3) 输出空间维度。批量或输出特征维度中将有一个被填充为 8 的倍数,而另一个被填充为 128 的倍数。输出空间维度不会被填充。

通常,对于具有空间维度(tf.nn.pooltf.conv2d 等)的操作,永远不会填充空间维度。

经验法则:为批量和特征维度选择高效的值

由于批量和特征维度可以填充,因此要谨慎确定批量和特征大小。为了充分利用 128 x 128 矩阵单元,请尽量为批量或特征(最好同时为这两者)设置合理大的值 (>=128)。

大多数情况下,批量大小为 128 足以保证完全占用 Cloud TPU 矩阵单元。Cloud TPU 上也支持使用较大批量大小运行,但建议使用 128 倍数的批量大小。如果您的模型无法在此类配置下运行,请尝试使用 8 的倍数的批量大小,从而最大限度地减少填充影响。

同样,特征维度也会映射到填充为 8 或 128 的维度,具体取决于 XLA 布局算法。也就是说,特征维度在是 8 或 128 的倍数时浪费的空间量最少。

融合

融合是 XLA 编译器用于优化程序的通用技术。 融合操作是指合并多个需要联合执行的操作。

例如,请考虑以下系列操作:

tmp = tf.add(x, y)
result = tf.multiply(tmp, z)

此代码大致相当于以下伪代码:

for (i = 0; i < element_count; i++) {
  tmp[i] = x[i] + y[i];
}

for (i = 0; i < element_count; i++) {
  result = tmp[i] * z[i];
}

通过融合,数组访问将同时进行:

for (i = 0; i < element_count; i++) {
  result = (x[i] + y[i]) * z[i];
}

在此例中,内存往返次数减少,XLA 无需为“tmp”分配任何空间。

融合是一项关键优化措施,可从多方面对 Cloud TPU 形成有利影响:

  • 无需在主内存中存储中间结果(通常很慢),可减少内存传输。
  • 可以更充分地利用硬件设备,以免这些设备被闲置。
  • 由于同时需要的活动缓冲区减少,可以降低模型的内存利用量。

广播

当两个形状不同但可兼容的张量合并时,系统将以隐式方式进行广播。

例如,tf.add(vector, matrix) 要求将矢量广播到矩阵的形状。操作结果具有与矩阵相同的形状。如需了解详情,请参阅广播数组指南。

虽然广播通常可与其使用方融合,但强制执行广播以实现物化会导致性能低下,同时增加内存压力。

在以下示例中,矢量与矩阵的相加操作中隐含的广播不能与 argmax 融合,因而导致广播物化:

tf.argmax(tf.add(vector, zero_matrix), axis=0)

TensorFlow 函数具体建议

请参阅 Cloud TPU 上可用的 TensorFlow 操作完整列表。

tf.matmul

  • 转置任一操作数的结果实际上是免费的。
  • 请注意,tf.matmul 支持融合其输入和输出。这意味着直接应用于 tf.matmul 输出的激活函数或偏差开销较低。

tf.nn.conv_n_dtf.nn.depthwise_conv2dtf.nn.separable_conv2d

  • 在激活时,批量和特征维度将填充为 8 或 128 的倍数。
    • 首先,XLA 会跟踪模块卷积中的最常见的批量维度大小。这有助于区分前向卷积、激活梯度卷积和内核梯度卷积。
    • 如果最常见的批量大小大于或等于 64:
      • 对于前向和后向卷积,将批量填充为 128 的倍数,而特征填充为 8 的倍数。
      • 对于梯度更新卷积,将批量填充为 8 的倍数,而特征填充为 128 的倍数。
    • 如果最常见的批量大小小于 64:
      • 对于前向和后向卷积,将批量填充为 8 的倍数,而特征填充为 128 的倍数。
      • 对于梯度更新卷积,将批量填充为 128 的倍数,而特征填充为 8 的倍数。
      • 如果转置仅交换输入特征和批量维度,则在将激活发送到卷积之前立即进行转置是免费的。
  • 对于内核,输入特征和输出特征维度将填充为 8 或 128 的倍数。确切的决定受内核的提供方和其他使用方的影响。
    • 如果转置仅交换输入和输出特征维度,则在将内核发送到卷积之前立即进行转置是免费的。
  • 对于结果,批量和特征维度将填充为 8 或 128 的倍数。
    • 如果转置仅交换批量和输出特征维度,则转置卷积结果是免费的。
  • 请注意,tf.nn.conv_n_d 支持融合其结果、激活和/或内核。这意味着直接应用于输出的激活函数或偏差开销较低。

tf.nn.avg_pooltf.nn.max_pool

  • 填充规则在此也适用:空间维度比批量和特征更重要。每个批量和特征可能被填充为 8 或 128 的倍数。
  • 通常,池操作的布局与流入或流出池的卷积匹配。
  • tf.nn.max_pool 的梯度计算可能比与其等效的 tf.nn.avg_pool 慢。请考虑尽可能从最大池化切换为平均池化。

tf.concattf.slicetf.strided_slice

  • 避免不必要的切片和串联。已填充维度中的切片和串联开销相当大。
    • 如果切片维度没有填充开销,则会最大程度减少数据移动。

tf.transpose

  • 可免费转置 tf.matmul 的任何操作数或结果。
  • 如果转置交换批量和输入特征维度,则将 tf.conv_n_d 的激活进行转置是免费的。
  • 如果转置交换输入和输出特征维度,则将 tf.conv_n_d 的内核进行转置是免费的。
  • 如果转置交换批量和输出特征维度,则将 tf.conv_n_d 的结果进行转置是免费的。

tf.batch_to_spacetf.space_to_batchtf.space_to_depthtf.depth_to_space

  • 由于这些操作需要在已填充的维度和未填充的维度间移动数据,因此成本很高。

tf.reshape

  • 在已填充的维度中移动数据时,Cloud TPU 上的重塑成本可能很高。
  • 如果存在大量填充,则在主机上将数据形状调整到 R1 并在设备上将其调整回维度更高的形状,可能会提高性能。这可以提高主机与设备之间的传输效率。
    • 由于这可以根据需要解压打包的参数,因而有助于降低峰值内存利用量。

tf.random_uniformtf.distributions.Bernoullitf.random_normaltf.multinomial

  • 用于实现均匀分布或伯努利分布的伪随机数生成非常快。
  • 正态分布比均匀分布或伯努利分布的开销略大。
  • 用于分类/多项分布的伪随机数生成开销相当大。

tf.reduce_alltf.reduce_anytf.reduce_logsumexptf.reduce_maxtf.reduce_mintf.reduce_prodtf.reduce_sum

  • 具有相同输入和输出形状的多次缩减可以通过融合并行执行。
    • 如果可能,尝试将连续的归约链重写为并行链。
  • 归约支持按元素将操作融入输入,但不支持融入输出。如果可能,请重写表达式以促进融合。 例如:

        tf.multiply(tf.reduce_sum(x), y)
    

    重写为:

        tf.reduce_sum(tf.multiply(x, y))
    

tf.nn.batch_normalizationtf.nn.fused_batch_normtf.layers.batch_normalization

  • Cloud TPU 编译器可以有减少低批量归一化的 TensorFlow 融合变体。它们比使用替代方案更有效。
    • 首选 tf.nn.fused_batch_norm,而不是 tf.nn.batch_normalization
    • 对于 tf.layers.batch_normalization,请将“fused”参数设置为 true。

tf.nn.depthwise_conv2dtf.nn.separable_conv2d

  • 这些参数尚未完全优化,性能可能不及其等效的常规卷积。