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

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

安装 TensorBoard

TensorBoard 会作为 TensorFlow 的一部分默认安装在 Cloud TPU 虚拟机中。您也可以手动安装 TensorFlow。 无论哪种情况,都可能需要一些额外的依赖项。通过运行以下命令进行安装:

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

安装 Cloud TPU TensorBoard 插件

通过 SSH 连接到您的 TPU 虚拟机:

   $ gcloud alpha compute tpus tpu-vm ssh your-vm --zone=your-zone
   

然后运行以下命令:

  pip3 install --upgrade "cloud-tpu-profiler>=2.3.0"
  pip3 install --user --upgrade -U "tensorboard>=2.3"
  pip3 install --user --upgrade -U "tensorflow>=2.3"

启动 TensorFlow Profiler 服务器

在开始训练循环之前,将以下代码添加到您的脚本中。

tf.profiler.experimental.server.start(6000)

这将在 TPU 虚拟机上启动 TensorFlow 性能剖析器服务器。

启动训练脚本

运行训练脚本并等待,直到您看到表明您的模型正在训练的输出。其样式取决于您的代码和模型。查找 Epoch 1/100 之类的输出。或者,您也可以导航到 GCP Console 中的 Cloud TPU 页面,选择您的 TPU,然后查看 CPU 利用率图表。虽然这并不会显示 TPU 利用率,但很好地表明 TPU 正在训练您的模型。

启动 TensorBoard 服务器

打开一个新的终端窗口,并使用端口转发通过 SSH 连接到您的 TPU 虚拟机。这样,您的本地浏览器就可以与 TPU 虚拟机上运行的 TensorBoard 服务器通信。

 gcloud alpha compute tpus tpu-vm ssh your-vm --zone=us-central1-a --ssh-flag="-4 -L 9001:localhost:9001"
 

在刚刚打开的终端窗口中运行 TensorBoard,并使用 --logdir 标志指定 TensorBoard 在其中写入分析数据的目录。例如:

TPU_LOAD_LIBRARY=0 tensorboard --logdir your-model-dir --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:6006/ (Press CTRL+C to quit)

打开网络浏览器并转到 TensorBoard 输出中显示的网址。从 TensorBoard 页面右上角的下拉菜单中选择 Profile

映像

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

  1. 选择“捕获性能剖析文件”按钮
  2. 选择“TPU 名称”单选按钮
  3. 输入您的 TPU 名称
  4. 选择“拍摄”按钮

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

捕获性能剖析文件后,TensorBoard 将显示 overview_page。在左侧窗格的工具下,选择 trace_viewer

映像

TPU 虚拟机目前不支持工具菜单中列出的其他工具。如需详细了解这些工具,请参阅在 TPU 节点上剖析模型

Trace Viewer

Trace viewer 是性能剖析 (Profile) 下的 Cloud TPU 性能分析工具。此工具使用 Chrome 跟踪记录事件分析查看器,因此仅适用于 Chrome 浏览器。

Trace Viewer 包含一个时间轴,其中显示的内容包括:

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

您可以通过 Trace Viewer 找出模型中的性能问题,然后采取措施来加以解决。例如,从总括层面来讲,您可以确定占用大部分时间的是馈入还是模型训练。展开细目,您可以确定哪些 TensorFlow 操作执行时间最长。

请注意,对于每个 Cloud TPU,Trace Viewer 最多限制 100 万个事件。如果您需要评估更多事件,请改用流式跟踪记录查看器

Trace Viewer 界面

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

映像

此屏幕包含以下主要元素(标有上述数字):

  1. 运行下拉列表。包含已捕获其跟踪记录信息的所有运行。默认视图为最近的运行,您也可以打开下拉列表选择其他运行。
  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 操作。显示在 TPU 上执行的 TensorFlow 操作。
    • XLA 操作。显示在 TPU 上运行的 XLA 操作。(每个操作都会转换为一个或多个 XLA 操作。XLA 编译器会将这些 XLA 操作转换为在 TPU 上运行的代码。)
  • 针对在主机的 CPU 上运行的线程的版块,其标签为“Host Threads”。对于每个 CPU 线程,此版块都包含一个跟踪记录。注意:您可以忽略版块标签旁边显示的信息。

时间轴工具选择器

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

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

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

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

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

注意:缩放工具有一个已知错误:当鼠标光标在时间轴视图外面时松开鼠标按钮,仍会进行缩放。如果您遇到这种情况,点击时间轴视图即可停止缩放。

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

请注意,如果您选择其他某个工具,间隔标记仍会保留。