使用 Cloud TPU 工具剖析模型的性能

重要提示 ：有两种 TPU 架构：TPU 虚拟机和 TPU 节点。这两种架构的性能剖析性能有所不同。本文档介绍了如何在 TPU 虚拟机上分析模型。系统架构中介绍了这两种 TPU 架构。

通过性能剖析，您可以优化模型在 Cloud TPU 上的训练性能。您可以使用 TensorBoard 和 Cloud TPU TensorBoard 插件来分析模型。

如需详细了解如何将 TensorBoard 与其中一个支持的框架搭配使用，请参阅以下文档：

• TensorFlow 性能指南
• PyTorch 性能指南
• JAX 性能指南

剖析训练脚本的前提条件

在使用 TPU 分析工具之前，您需要：

1. 启动模型训练会话

   1. 设置 v4-8 TPU 来训练模型。本文档中介绍的分析过程使用 ResNet 模型，但您也可以使用另一个模型，前提是该模型在 v4 TPU 上训练。

   2. 在 TPU 虚拟机中，添加一行代码来启动性能分析器服务器到训练脚本。

      对于 ResNET 训练，训练脚本位于：/usr/share/tpu/tensorflow/resnet50_keras/resnet50.py。

      将突出显示的代码行插入 resnet50.py 中。 在文件顶部，添加以下导入操作：

      import tensorflow.compat.v2 as tf2
      

      在脚本开始训练循环之前，添加突出显示的代码行：

      if name == 'main':
       tf.logging.set_verbosity(tf.logging.INFO)
       tf2.profiler.experimental.server.start(6000)
       app.run(main)
      

      当您运行脚本时，TensorFlow 性能分析器服务器会从您的 TPU 虚拟机启动。

   3. 开始模型训练。

      运行训练脚本，然后等到看到输出表明模型正在主动训练时。输出取决于您的代码和模型。查找类似 Epoch 1/100 的输出。 或者，您也可以导航到 Google Cloud 控制台中的 Cloud TPU 页面、选择您的 TPU，然后查看 CPU 利用率图表。虽然 CPU 利用率图未显示 TPU 利用率，但充分表明 TPU 正在训练模型。

开始分析模型训练

在模型训练期间，打开一个单独的终端窗口或 Cloud Shell。按照以下步骤开始分析模型训练。

1. 在新的窗口或 shell 中，通过端口转发连接到 TPU 虚拟机。

   gcloud compute tpus tpu-vm ssh your-vm --zone=us-central2-b --ssh-flag="-4 -L 9001:localhost:9001"
   

   端口转发允许本地浏览器与 TPU 虚拟机上运行的 TensorBoard 服务器通信。

2. 安装 TensorFlow 要求 {: id="install-Tensorboard"}。

   您的 TPU 虚拟机默认安装了 TensorBoard。您也可以手动安装 TensorFlow。 无论采用哪种方式，都可能需要一些额外的依赖项。请运行以下命令，在您的 TPU 虚拟机上安装这些依赖项：

   pip3 install -r /usr/share/tpu/models/official/requirements.txt
   

3. 安装 Cloud TPU TensorBoard 插件 {: id="install-plugin"}。

   从 TPU 虚拟机运行以下命令：

    pip3 install --upgrade "cloud-tpu-profiler>=2.3.0"
    pip3 install tensorflow
    pip3 install tensorboard_plugin_profile
   

4. 启动 TensorBoard 服务器

   运行 TensorBoard 并在 TPU 虚拟机上创建一个日志目录 (logdir)，TensorBoard 虚拟机上会写入 TensorBoard 分析数据。使用 --logdir 标志指定日志目录。例如：

   mkdir log-directory
   TPU_LOAD_LIBRARY=0 tensorboard --logdir log-directory --port 9001
   

TensorBoard 会启动一个网络服务器并显示其网址：

Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.3.0 at http://localhost:9001 (Press CTRL+C to quit)

打开网络浏览器并转到 TensorBoard 输出中显示的网址。从 TensorBoard 页面右上角的下拉菜单中选择分析。可用性能剖析工具的列表显示在左侧边栏中的工具下拉菜单中。

捕获 TPU 虚拟机上的性能剖析文件

1. 选择拍摄个人资料按钮。
2. 选择 IP 地址单选按钮。
3. 在 Profile Service URL 字段中输入 HostNAME:6000。
4. 选择拍摄按钮。

使用 TensorBoard 查看性能剖析文件数据

捕获性能剖析文件后，TensorBoard 会显示 overview_page。您可以使用的分析工具列表显示在左侧窗格中。

个人资料

在您捕获某些模型数据后，系统将显示 Profile 标签页。您可能需要点击 TensorBoard 页面上的刷新按钮。数据可用后，点击分析标签页会显示一系列有助于性能分析的工具。您可以使用以下任一工具来分析模型。

• 概览页面
• Input Pipeline Analyzer
• XLA Op Profile
• Trace Viewer（仅限 Chrome 浏览器）
• Memory Viewer

性能剖析概览页面

“分析”页面中的“概览”页面 (overview_page) 提供了模型在捕获运行期间的执行情况汇总。该页面显示了所有 TPU 的聚合概览以及整体输入流水线分析。您可以在主机下拉列表中选择单个 TPU。

该页面在以下面板中显示数据：

• 性能摘要

  • FLOPS 利用率 - TPU 矩阵单元的使用百分比

• TPU 上的十大 TensorFlow 操作显示消耗最多的 TensorFlow 操作：

  每行显示某项操作的自用时间（以所有操作所用时间的百分比的形式）、累计时间、类别、名称以及实现的 FLOPS 率。

• 运行环境

  • 使用的主机数量
  • 使用的 TPU 类型
  • TPU 核心数量

Input Pipeline Analyzer

Input Pipeline Analyzer 可提供有关性能结果的数据分析。该工具展示 capture_tpu_profile 工具收集的 input_pipeline.json 文件的性能结果。

该工具能够立即告诉您程序是否输入受限，并可指导您完成设备端和主机端分析，以调试流水线上任何存在瓶颈的阶段。

如需深入了解如何优化流水线性能，请参阅有关输入流水线性能的指南。

输入流水线

当 TensorFlow 程序从文件中读取数据时，读取过程会分为多个串联连接的数据处理阶段。一个阶段的输出是下一个阶段的输入。这个读取系统称为输入流水线。

从文件中读取记录的典型流水线具有以下阶段：

1. 文件读取
2. 文件预处理（可选）
3. 将文件从宿主机传输到设备

效率低下的输入流水线可能会严重降低应用的速度。当应用在其输入流水线上花费大量时间时，我们将其称为输入受限。使用输入流水线分析器可以了解输入流水线的效率低下。

输入流水线信息中心

如需打开输入流水线分析器，请选择 Profile，然后从工具下拉列表中选择 input_pipeline_analyzer。

信息中心会显示设备端和主机端分析的详细信息。

设备端分析 - 显示设备单步用时的详细信息。

• 设备单步用时统计信息
• 设备单步用时等待输入数据的百分比

主机端分析

本部分详细介绍了分为几类主机端分析：

• 将要传输到设备的数据加入队列。向设备传输数据之前将数据加入馈入队列所花费的时间。
• 数据预处理：在预处理操作（例如图像解压缩）上所花的时间。
• 预先从文件中读取数据：读取文件所用的时间，包括缓存、预取和交错。
• 按需从文件中读取数据：在没有缓存、预取和交错的情况下从文件读取数据所用的时间。
• 其他数据读取或处理：在未使用 tf.data 的其他输入相关操作上花费的时间。

如需查看按执行时间细分的单个输入操作及其类别的统计信息，请展开 Show Input Op statistics 部分。

系统将显示如下所示的源数据表：

每个表条目包含以下信息：

1. 输入操作显示输入操作的 TensorFlow 操作名称。
2. 计数：显示分析期间执行的操作的实例总数。
3. 总时间（以毫秒为单位）显示每个操作实例的累计时间总和。
4. 总时间百分比显示一项操作所用的总时间占输入处理总时间的比例。
5. 总自用时间（以毫秒为单位）显示函数的所有实例的累计时间。自用时间测量在函数正文内部花费的时间，不包括在它调用的任何函数上花费的时间。例如，Iterator::PaddedBatch::Filter::ForeverRepeat::Map 由 Iterator::PaddedBatch::Filter 调用，因此前者的总自用时间不包括在后者的总自用时间内。
6. 总自用时间百分比显示总自用时间占输入处理总时间的比例。
7. Category 显示输入操作的处理类别。

Op Profile

Op Profile () 是一款 Cloud TPU 工具，它可展示在性能剖析期间执行的 XLA 操作的性能统计信息。操作配置文件显示：

• 您的应用使用 Cloud TPU 的程度，按类别和 TPU FLOPS 利用率显示各操作所用时间的百分比。
• 最耗时的操作。这些操作可能有待优化。
• 各个操作的详细信息，包括形状、填充和使用该操作的表达式。

您可以使用操作配置文件查找优化目标。例如，您可以使用操作配置文件来确定哪些 XLA 操作运行时间最长以及它们使用了多少 TPU FLOPS。

使用 Op Profile

在性能剖析文件收集期间，capture_tpu_profile 还会创建一个包含 XLA 操作的性能统计信息的 op_profile.json 文件。

您可以在 TensorBoard 中查看 op_profile 中的数据，方法是：点击屏幕顶部的 Profile 标签页，然后从 Tools 下拉列表中选择 op_profile。您将看到如下所示的界面：

1. “概览”部分显示了 Cloud TPU 的利用率并提供了优化建议。
2. 控制台 包含一些控件，通过这些控件，您可以设置在表中显示的操作数量、显示哪些操作，以及对这些操作进行排序的方式。
3. 操作表：列出与 XLA 操作相关联的热门 TensorFlow 操作类别的表。这些操作按 Cloud TPU 使用率百分比排序。
4. 操作详细信息卡片显示指向表中的操作时显示的操作的详细信息。这些详细信息包括 FLOPS 利用率、操作使用的表达式以及操作布局（适合）。

XLA 操作表

操作表按照 Cloud TPU 用量百分比从高到低的顺序列出 XLA 操作类别。该表显示了所用时间的百分比、操作类别名称、关联的 TensorFlow 操作名称，以及该类别的 FLOPS 利用率百分比。如需显示（或隐藏）某个类别的 10 个最耗时的 XLA 操作，请点击表格中类别名称旁边的三角形。

1. 时间显示该类别中所有操作所用时间的总百分比。您可以点击以展开条目，查看每个操作所用时间的明细。
2. 十大操作：类别名称旁边的切换开关可显示/隐藏类别中排名前十的耗时操作。如果操作列表中显示了融合操作条目，您可以展开该条目以查看其包含的非融合元素智能操作。
3. TensorFlow 操作显示与 XLA 操作关联的 TensorFlow 操作名称。
4. FLOPS 显示 FLOPS 利用率，即以 Cloud TPU 峰值 FLOPS 的百分比表示的 FLOPS 的测量值。FLOPS 利用率百分比越高，操作运行速度越快。表单元格采用颜色编码：绿色表示 FLOPS 利用率高（好），红色表示 FLOPS 利用率低（差）。

操作详细信息卡片

选择表条目时，系统会显示一张卡片，其中显示了有关 XLA 操作或操作类别的详细信息。典型的卡片如下所示：

• Name 和 Category 显示突出显示的 XLA 操作名称和类别。
• FLOPS 利用率显示 FLOPS 利用率，以占总 FLOPS 的百分比表示。
• 表达式显示包含该操作的 XLA 表达式。
• 内存利用率显示程序峰值内存用量的百分比。
• 布局（仅限卷积运算）显示张量的形状和布局，包括对 XLA 编译器执行的任何内边距的说明。

解读结果

对于卷积运算，TPU FLOPS 利用率低可能是由以下一个或两个原因造成的：

• 填充（矩阵单元仅被部分使用）
• 卷积运算受限于内存

本部分解释了 FLOP 利用率较低的模型中的一些性能指标。在此示例中，执行时间占主导地位，输出融合和卷积占据了主导地位。有许多矢量或标量运算的 FLOP 利用率较低。

此类性能剖析文件的一种优化策略是将矢量或标量操作转换为卷积运算。

在以下示例中，%convolution.399 显示的 FLOPS 利用率和内存利用率低于上例中的 %convolution.340。

在此示例中，批量大小填充到 128，特征大小填充到 8。在本例中，只有 5% 的矩阵单元有效使用。利用率的计算方法是：(((time_time * num_of_features) / padding_size ) / num_of_cores)。 将此示例中的 FLOPS 与上一个示例中不使用内边距的 %convolution.340 进行比较。

Trace Viewer

Trace Viewer 是配置文件页面中提供的 Cloud TPU 性能分析工具。此工具使用 Chrome 跟踪记录事件分析查看器，因此仅适用于 Chrome 浏览器。

Trace Viewer 包含一个时间轴，其中显示了以下内容：

• 由 TensorFlow 模型执行的操作的时长。
• 系统的哪个部分（TPU 还是宿主机）执行了操作。通常，宿主机执行馈入操作来预处理训练数据并将其传输到 TPU，而 TPU 执行实际模型训练。

Trace Viewer 可让您发现模型中的性能问题，然后采取措施来修复这些问题。例如，在较高层面上，您可以确定是信息流还是模型训练花费的时间最多。展开细目，您可以确定哪些 TensorFlow 操作执行时间最长。

Trace Viewer 针对每个 Cloud TPU 仅限访问 100 万个事件。如果您需要评估更多事件，请改用流式跟踪记录查看器。

Trace Viewer 界面

如需打开 Trace Viewer，请转到 TensorBoard，点击屏幕顶部的性能剖析 (Profile) 标签页，然后从工具 (Tools) 下拉列表中选择 trace_viewer。Trace Viewer 会显示最近的运行：

此屏幕包含以下主要元素（在上一个屏幕截图中标有数字）：

1. Runs 下拉菜单 包含您已捕获其跟踪记录信息的所有运行。默认视图是最近一次运行，但您可以打开下拉菜单选择其他运行。
2. 工具下拉菜单：选择不同的分析工具。
3. 主机下拉菜单：选择包含 Cloud TPU 集的主机。
4. “时间轴”窗格显示 Cloud TPU 和宿主机在一段时间内执行的操作。
5. 详细信息窗格：显示“时间轴”窗格中所选操作的更多信息。

以下是时间轴窗格的详细信息：

时间轴窗格包含以下元素：

1. 顶部栏 包含各种辅助控件。
2. 时间轴显示相对于跟踪记录开始时的时间。
3. 版块和轨道标签：每个版块包含多个轨道，左侧有一个三角形，点击该按钮即可展开和收起。系统中的每个处理元素都有一个版块。
4. 工具选择器 包含与 Trace Viewer 交互的各种工具。
5. 事件显示操作的执行时间或元事件（例如训练步）的时长。
6. 垂直标签页栏：此栏对于 Cloud TPU 没有用处。竖条是 Chrome 提供的通用跟踪记录查看器工具的一部分，用于各种性能分析任务。

版块和跟踪记录

Trace Viewer 包含以下版块：

• 每个 TPU 节点有一个版块，上面标有 TPU 芯片的编号及芯片内 TPU 节点的编号（例如“芯片 2: TPU 核心 1”）。每个 TPU 节点版块包含以下跟踪记录：
  • 步骤：显示 TPU 上运行的训练步骤的时长。
  • TensorFlow Ops 显示在 TPU 上执行的 TensorFlow 操作。
  • XLA Ops 显示在 TPU 上运行的 XLA 操作。（每个操作都会转换为一个或多个 XLA 操作。XLA 编译器会将这些 XLA 操作转换为在 TPU 上运行的代码。）
• 针对在主机 CPU 上运行的线程的版块，其标签为 "Host Threads"。对于每个 CPU 线程，此版块都包含一个跟踪记录。注意：您可以忽略版块标签旁边显示的信息。

时间轴工具选择器

您可以使用 TensorBoard 中的时间轴工具选择器与时间轴视图进行交互。您可以点击时间轴工具或使用以下键盘快捷键激活并突出显示该工具。如需移动时间轴工具选择器，请点击顶部的虚线区域，然后将选择器拖动到所需位置。

以下是各个时间轴工具的使用方法：

	选择工具 点击某个事件即可将其选中，拖动即可选择多个事件。有关所选事件（名称、开始时间和时长）的更多信息将显示在“详细信息”窗格中。
	平移工具 拖动可在水平和垂直方向平移时间轴视图。
	缩放工具 沿水平（时间）轴方向，向上拖动可放大视图，向下拖动可缩小视图。鼠标光标的水平位置决定缩放操作的中心。 注意：如果在松开鼠标按钮后缩放工具保持活动状态，请点击时间轴视图以停用缩放工具。
	计时工具 水平拖动可以标记时间间隔。间隔的时长会显示在时间轴上。如需调整间隔，请拖动其末端。如需清除间隔，可点击时间轴视图内的任意位置。 如果您选择其他工具，间隔值仍会保持标记状态。

Memory Viewer

使用 Memory Viewer，您可以直观了解程序的峰值内存用量和内存用量趋势。

Memory Viewer 界面如下所示：

1. Host 主机：选择一个 TPU 主机和 XLA 高级优化器 (HLO) 模块进行可视化。
2. 内存概览显示峰值内存分配和大小（不填充）。
3. 工作空间图表显示程序的峰值内存用量以及内存使用量的趋势。指向某个缓冲区图表中的缓冲区，即可在缓冲区分配卡片中显示更多信息。
4. 缓冲区图表：显示峰值内存用量时缓冲区分配的两个图表。指向某个缓冲区图表中的缓冲区，即可在缓冲区详细信息卡片中显示更多信息。
5. 缓冲区分配详细信息卡片：显示缓冲区的分配详细信息。

内存概览面板

内存概览（顶部）面板显示模块名称，以及总缓冲区分配大小达到最大值时的峰值内存分配设置。此面板还显示未填充的峰值分配大小以供比较。

工作空间图表

此图表显示峰值内存用量并绘制程序的内存用量趋势。垂直线表示该程序的峰值内存利用率。此图表显示了您的程序能否放到可用的全局内存空间中。

图中的每个点表示 XLA HLO 程序中的“程序点”。该行显示了程序的内存使用量随时间的变化情况。

与缓冲区图表元素的交互

当您指向某个缓冲区图中的缓冲区时，会在工作空间图表中显示一条显示该缓冲区生命周期的水平线。

水平线的粗细表示缓冲区大小相对于峰值内存分配的相对量级。线长表示缓冲区的生命周期。

缓冲区图表

有两个图表显示峰值用量下的内存用量明细。

• 按程序顺序按照缓冲区在程序执行期间活跃的顺序，从左到右显示缓冲区。

• By Size，按程序大小顺序显示程序执行期间有效的缓冲区。

缓冲区分配详细信息卡片

当您指向其中一个缓冲区图表中显示的缓冲区时，系统会显示缓冲区分配详细信息卡片。典型的详细信息卡片如下所示：

1. 名称 - XLA 操作的名称。
2. 类别 - 操作类别。
3. Size - 缓冲区分配的大小（包括填充）。
4. 未填充大小 - 无填充的缓冲区分配的大小。
5. 展开 - 填充的缓冲区大小与未填充的大小的相对量级。
6. Extra memory - 指示用于填充的内存量。
7. 形状 - 用于描述 N 维数组的阶、大小和数据类型。
8. TensorFlow 操作名称 - 显示与缓冲区分配关联的 TensorFlow 操作的名称。
9. Allocation type - 指示缓冲区分配类别：参数、输出、线程局部和临时（例如，融合中的缓冲区分配）。

内存不足错误

如果您运行模型并收到“内存不足错误”，请使用以下命令捕获内存配置文件。等到脚本训练模型后，再启动性能分析器。您可以使用 duration_ms 标志指定性能分析器的运行时长。此命令的输出可以帮助您了解出现错误的原因：

  (vm)$ capture_tpu_profile --tpu=$TPU_NAME --logdir=log-directory --duration_ms=60000