在 Cloud TPU 虚拟机上分析模型

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

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

分析训练脚本的前提条件

在使用 TPU 性能剖析工具之前,您需要:

  1. 启动模型训练会话

    1. 设置 v4-8 TPU 以训练模型。本文档中所述的分析过程使用 ResNet 模型,但您可以使用另一个模型,前提是该模型在 v4 TPU 上训练。
    2. 在 TPU 虚拟机中,向训练脚本添加一行以启动性能分析器服务器。

      对于 ResNET 训练,训练脚本位于:/usr/share/tpu/tensorflow/resnet50_keras/resnet50.py

      将突出显示的代码行插入 resnet50.py。 在文件顶部,添加以下 import:

      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 可以从中写入性能剖析数据。使用 --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 页面右上角的下拉菜单中选择分析 (Profile)。可用性能剖析工具的列表显示在左侧边栏的tools下拉菜单中。

图片

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

  1. 选择捕获配置文件按钮。
  2. 选择 IP 地址单选按钮。
  3. Profile Service URL 字段中输入 HOSTNAME:6000
  4. 选择拍摄按钮。

图片

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

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

图片

配置文件

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

性能剖析概览页面

配置文件页面中的概览页面 (overview_page) 提供模型在捕获运行期间的执行情况汇总。此页面会显示所有 TPU 的汇总概览以及整体输入流水线分析。您可以在主机下拉列表中选择各个 TPU。

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

图片

  • 性能摘要

    • FLOPS 利用率 - TPU 矩阵单元的使用百分比
  • TPU 上排名前 10 的 TensorFlow 操作显示消耗最多的 TensorFlow 操作:

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

  • 运行环境

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

Input Pipeline Analyzer

Input Pipeline Analyzer 可提供有关性能结果的数据分析。该工具可立即告知您程序是否受限于输入,并且可以引导您完成设备和主机端分析,以调试正在造成瓶颈的任何流水线阶段。

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

输入流水线

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

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

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

效率低下的输入流水线可能会严重降低应用的速度。当应用在其输入流水线上花费大量时间时,我们将其称为输入受限。使用 Input Pipeline Analyzer 了解输入流水线的低效环节。

输入流水线信息中心

如需打开 Input Pipeline Analyzer,请选择性能剖析 (Profile),然后从工具下拉列表中选择 input_pipeline_analyzer

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

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

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

主机端分析

本部分显示主机端分析的详细信息,分为以下几类:

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

图片

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

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

图片

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

  1. 输入操作 (Input Op) 显示输入操作的 TensorFlow 操作名称。
  2. 计数 显示在分析期间执行的操作的实例总数。
  3. 总时间(以毫秒为单位)显示每个操作实例所用时间的累计总和。
  4. 总时间百分比:显示一项操作上花费的总时间占输入处理总时间的比例。
  5. 总自用时间(以毫秒为单位)显示函数所有实例的累计时间。自用时间测量在函数体内所用的时间,不包括它调用的任何函数所用的时间。例如,Iterator::PaddedBatch::Filter::ForeverRepeat::MapIterator::PaddedBatch::Filter 调用,因此其总自用时间不会计入后者的总自用时间。
  6. 总自用时间百分比 显示总自用时间占输入处理总时间的比例。
  7. 类别:显示输入操作的处理类别。

Op Profile

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

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

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

使用 Op Profile

Op Profile 工具包含 XLA 操作的性能统计信息。您可以在 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. 前 10 项操作:类别名称旁边的切换开关可显示/隐藏该类别中 10 项最耗时的操作。如果操作列表中显示了融合操作条目,您可以展开该条目以查看其中包含的非融合元素级操作。
  3. TensorFlow 操作:显示与 XLA 操作关联的 TensorFlow 操作名称。
  4. FLOPS 显示 FLOPS 利用率,即以 Cloud TPU 峰值 FLOPS 的百分比表示的 FLOPS 的测量值。FLOPS 利用率百分比越高,操作运行速度越快。表单元格采用颜色编码:绿色表示 FLOPS 利用率高(好),红色表示 FLOPS 利用率低(差)。

操作详细信息卡片

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

图片

  • 名称类别:显示突出显示的 XLA 操作名称和类别。
  • FLOPS 利用率以可能的总 FLOPS 的百分比显示 FLOPS 利用率。
  • 表达式:显示包含相应操作的 XLA 表达式
  • 内存利用率显示程序的峰值内存使用量百分比。
  • 布局(仅限卷积运算)显示张量的形状和布局,包括由 XLA 编译器执行的任何填充的说明。

解读结果

对于卷积操作,TPU FLOPS 利用率低可能是由于以下一个或两个原因造成的:

  • 填充(矩阵单元仅被部分使用)
  • 卷积运算受内存限制

本部分对 FLOP 利用率低的模型的一些性能指标的解释。在此示例中,输出融合卷积占据了主要执行时间。许多矢量或标量操作的 FLOP 利用率较低。

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

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

图片

在此示例中,批次大小被填充为 128,特征大小被填充为 8。在这种情况下,只有 5% 的矩阵单元得到有效利用。利用率计算公式为:(((batch_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. Tools 下拉列表,用于选择不同的性能分析工具。
  3. 主机下拉列表:选择包含 Cloud TPU 集的主机。
  4. 时间轴窗格 显示 Cloud TPU 和宿主机随时间执行的操作。
  5. 详细信息窗格:显示时间轴窗格中所选操作的更多信息。

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

图片

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

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

版块和跟踪记录

Trace Viewer 包含以下版块:

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

时间轴工具选择器

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

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

选择工具
点击事件即可将其选中,拖动可选择多个事件。 有关所选事件(名称、开始时间和时长)的更多信息将显示在“详细信息”窗格中。

平移工具
拖动可在水平和垂直方向平移时间轴视图。

缩放工具
沿水平(时间)轴方向,向上拖动可放大视图,向下拖动可缩小视图。鼠标光标的水平位置决定缩放操作的中心。

注意:如果松开鼠标按钮后缩放工具仍处于活动状态,请点击时间轴视图以停用缩放工具。

计时工具
水平拖动可以标记时间间隔。间隔的时长会显示在时间轴上。如需调整间隔,请拖动其末端。如需清除间隔,可点击时间轴视图内的任意位置。

如果您选择其他工具,时间间隔仍会保留标记状态。

Memory Viewer

借助 Memory Viewer,您可以直观地了解程序的峰值内存用量和内存用量趋势。

Memory Viewer 界面如下所示:

图片

  1. 主机 下拉列表:选择要直观呈现的 TPU 主机和 XLA 高级优化器 (HLO) 模块。
  2. 内存概览显示峰值内存分配和无填充的大小。
  3. 工作空间图表:显示程序的峰值内存用量和内存用量趋势图。将光标指向某个缓冲区图表中的缓冲区,以在缓冲区分配卡中显示更多信息。
  4. 缓冲区图表:显示峰值内存使用时的缓冲区分配情况的两个图表。指向某个缓冲区图表中的缓冲区,以在缓冲区详细信息卡片中显示其他信息。
  5. 缓冲区分配详细信息卡片 - 显示缓冲区分配详细信息。

内存概览面板

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

图片

工作空间图表

此图表显示峰值内存使用率和程序的内存使用趋势图。垂直线表示程序的峰值内存利用率。此图表显示了您的程序是否可以放入可用的全局内存空间中。

图片

图表中的每个点代表 XLA HLO 程序中的“程序点”。此线条显示程序的内存用量随时间的变化情况。

与缓冲区图表元素的交互

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

图片

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

缓冲区图表

两个图表显示了峰值使用量时的内存使用量明细。

图片

  • 按程序顺序:按程序执行期间缓冲区的活动顺序显示缓冲区(从左到右)。

  • 按大小显示程序执行期间处于活动状态的缓冲区(按大小降序排列)。

缓冲区分配详细信息卡片

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

图片

  1. 名称 - XLA 操作的名称。
  2. 类别 - 操作类别。
  3. Size - 缓冲区分配的大小(包括填充)。
  4. 未填充的大小 - 无填充的缓冲区分配的大小。
  5. 扩展 - 填充的缓冲区空间与未填充的缓冲区空间的相对量级。
  6. 额外内存 - 表示填充使用了多少额外内存。
  7. 形状 - 描述 N 维数组的阶、大小和数据类型。
  8. TensorFlow 操作名称 - 显示与缓冲区分配关联的 TensorFlow 操作的名称。
  9. 分配类型 - 表示缓冲区分配类别:参数、输出、线程局部和临时(例如,融合中的缓冲区分配)。

内存不足错误

如果您在运行模型时收到“内存不足”错误,请按照本文档中的准则来捕获性能分析报告。等待脚本完成模型训练,然后再启动性能分析器。此性能剖析输出有助于您了解导致错误的原因。