分析 TensorFlow 工作负载

如需在 Cloud TPU 上分析 TensorFlow 模型，您可以使用 TensorBoard 和 TPU TensorBoard 插件。TensorBoard 预安装在 TPU 虚拟机上。如需了解如何安装 TPU TensorBoard 插件并捕获性能配置文件，请参阅使用 Cloud TPU 工具分析模型。对于 有关 Cloud TPU 性能的一般信息，请参阅 Cloud TPU 性能指南。

如需了解详情，请参阅 TensorBoard 回调。

TensorFlow 函数性能说明

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

tf.matmul

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

tf.nn.conv_n_d、tf.nn.depthwise_conv2d、tf.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_pool、tf.nn.max_pool

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

tf.concat、tf.slice、tf.strided_slice

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

tf.transpose

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

tf.batch_to_space、tf.space_to_batch、tf.space_to_depth、tf.depth_to_space

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

tf.reshape

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

tf.random_uniform、tf.distributions.Bernoulli、tf.random_normal、tf.multinomial

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

tf.reduce_all、tf.reduce_any、tf.reduce_logsumexp、tf.reduce_max、tf.reduce_min、tf.reduce_prod、tf.reduce_sum

  • 可以在 并行处理。
    • 如果可能，尝试将连续的归约链重写为并行链。

  • 归约支持按元素将操作融入输入，但不支持融入输出。如果可能，请重写表达式以促进融合。 例如：

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

    重写为：

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

tf.nn.batch_normalization、tf.nn.fused_batch_norm、tf.layers.batch_normalization

  • XLA 编译器可以高效地降低 TensorFlow 的融合变体 和批量归一化的过程。它们比使用替代方案更有效。

    • 首选 tf.nn.fused_batch_norm，而不是 tf.nn.batch_normalization。
    • 对于 tf.layers.batch_normalization，请将“fused”参数设置为 true。