分析 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。