计算分析

计算分析#

使用 Perfetto 查看程序跟踪#

我们可以使用 JAX 分析器来生成 JAX 程序的跟踪，这些跟踪可以使用 Perfetto 可视化工具 进行可视化。目前，这种方法会阻塞程序，直到点击链接并且 Perfetto UI 加载跟踪。如果你希望在不进行任何交互的情况下获取分析信息，请查看下面的 Tensorboard 分析器。

with jax.profiler.trace("/tmp/jax-trace", create_perfetto_link=True):
  # Run the operations to be profiled
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

完成此计算后，程序将提示您打开一个指向 ui.perfetto.dev 的链接。当您打开链接时，Perfetto UI 将加载跟踪文件并打开一个可视化工具。

Perfetto 跟踪查看器

加载链接后，程序将继续执行。链接在打开一次后将不再有效，但它会重定向到一个新的有效URL。然后，您可以在Perfetto UI中点击“分享”按钮，创建一个可以与他人共享的跟踪记录的固定链接。

远程分析#

当分析运行在远程（例如在托管的虚拟机上）的代码时，你需要在端口9001上建立一个SSH隧道以使链接工作。你可以使用以下命令来完成：

$ ssh -L 9001:127.0.0.1:9001 <user>@<host>

或者，如果您使用的是 Google Cloud：

$ gcloud compute ssh <machine-name> -- -L 9001:127.0.0.1:9001

手动捕获#

除了使用 jax.profiler.trace 以编程方式捕获跟踪外，您还可以通过调用 jax.profiler.start_server(<port>) 在感兴趣的脚本中启动一个分析服务器。如果您只需要分析服务器在脚本的一部分中处于活动状态，您可以通过调用 jax.profiler.stop_server() 来关闭它。

一旦脚本运行并且分析服务器启动后，我们可以通过运行以下命令手动捕获和追踪：

$ python -m jax.collect_profile <port> <duration_in_ms>

默认情况下，生成的跟踪信息会转储到一个临时目录中，但可以通过传入 --log_dir=<选择的目录> 来覆盖此设置。此外，默认情况下，程序会提示您打开一个指向 ui.perfetto.dev 的链接。当您打开链接时，Perfetto UI 将加载跟踪文件并打开一个可视化工具。通过在命令中传入 --no_perfetto_link 可以禁用此功能。或者，您也可以将 Tensorboard 指向 log_dir 以分析跟踪信息（参见下面的“Tensorboard 分析”部分）。

TensorBoard 分析#

TensorBoard 的分析器 可以用于分析 JAX 程序。TensorBoard 是一个获取和可视化程序性能跟踪和分析的好方法，包括 GPU 和 TPU 上的活动。最终结果看起来像这样：

TensorBoard 分析器示例

安装#

TensorBoard 分析器仅在使用 TensorFlow 捆绑版本的 TensorBoard 时可用。

pip install tensorflow tensorboard-plugin-profile

如果你已经安装了 TensorFlow，你只需要安装 tensorboard-plugin-profile pip 包。注意只安装一个版本的 TensorFlow 或 TensorBoard，否则你可能会遇到 下面 描述的“重复插件”错误。更多关于安装 TensorBoard 的信息，请参见 https://www.tensorflow.org/guide/profiler。

程序化捕获#

你可以通过 jax.profiler.start_trace() 和 jax.profiler.stop_trace() 方法来记录代码的性能分析跟踪。使用要写入跟踪文件的目录调用 start_trace()。这应该是用于启动 TensorBoard 的相同 --logdir 目录。然后，你可以使用 TensorBoard 来查看跟踪结果。

例如，要进行性能分析跟踪：

import jax

jax.profiler.start_trace("/tmp/tensorboard")

# Run the operations to be profiled
key = jax.random.key(0)
x = jax.random.normal(key, (5000, 5000))
y = x @ x
y.block_until_ready()

jax.profiler.stop_trace()

注意 block_until_ready() 调用。我们使用这个来确保设备上的执行被跟踪捕获。有关为什么需要这个的详细信息，请参见 异步分发。

你也可以使用 jax.profiler.trace() 上下文管理器作为 start_trace 和 stop_trace 的替代方案：

import jax

with jax.profiler.trace("/tmp/tensorboard"):
  key = jax.random.key(0)
  x = jax.random.normal(key, (5000, 5000))
  y = x @ x
  y.block_until_ready()

要查看跟踪，如果尚未启动，请先启动 TensorBoard：

$ tensorboard --logdir=/tmp/tensorboard
[...]
Serving TensorBoard on localhost; to expose to the network, use a proxy or pass --bind_all
TensorBoard 2.5.0 at http://localhost:6006/ (Press CTRL+C to quit)

在这个例子中，你应该能够在 http://localhost:6006/ 加载 TensorBoard。你可以使用 --port 标志指定不同的端口。如果在远程服务器上运行 JAX，请参见下面的 远程分析。

然后，可以在右上角的“Profile”下拉菜单中选择“Profile”，或者直接访问 http://localhost:6006/#profile。左侧的“Runs”下拉菜单中会显示可用的跟踪记录。选择你感兴趣的运行记录，然后在“Tools”下选择 trace_viewer。你现在应该会看到执行的时间线。你可以使用 WASD 键来导航跟踪记录，点击或拖动以选择事件以查看底部的更多详细信息。有关使用跟踪查看器的更多详细信息，请参阅 这些 TensorFlow 文档。

你也可以使用 memory_viewer、op_profile 和 graph_viewer 工具。

通过 TensorBoard 手动捕获#

以下是从运行程序中捕获手动触发的 N 秒跟踪的说明。

  1. 启动一个 TensorBoard 服务器：

    tensorboard --logdir /tmp/tensorboard/

    你应该能够在 http://localhost:6006/ 加载 TensorBoard。你可以使用 --port 标志指定不同的端口。如果在远程服务器上运行 JAX，请参见下面的 远程分析。

  2. 在您想要分析的 Python 程序或进程中，在开始附近添加以下内容：

    import jax.profiler
jax.profiler.start_server(9999)

    这将启动 TensorBoard 连接的分析器服务器。在进入下一步之前，分析器服务器必须正在运行。当您使用完服务器后，可以调用 jax.profiler.stop_server() 来关闭它。

    如果你想分析一个长时间运行的程序片段（例如一个长时间的训练循环），你可以在程序的开头放置这个，并像往常一样启动你的程序。如果你想分析一个短程序（例如一个微基准测试），一种选择是在IPython shell中启动分析器服务器，并在下一步启动捕获后使用%run运行短程序。另一种选择是在程序开始时启动分析器服务器，并使用time.sleep()给你足够的时间来启动捕获。

  3. 打开 http://localhost:6006/#profile，然后点击左上角的“CAPTURE PROFILE”按钮。输入“localhost:9999”作为配置文件服务URL（这是您在上一步中启动的分析服务器地址）。输入您希望分析的毫秒数，然后点击“CAPTURE”。

  4. 如果你想要分析的代码还没有运行（例如，如果你在Python shell中启动了分析服务器），请在捕获运行时运行它。

  5. 捕获完成后，TensorBoard 应该会自动刷新。（并非所有的 TensorBoard 分析功能都与 JAX 挂钩，因此最初可能会看起来没有捕获到任何内容。）在左侧的“工具”下，选择 trace_viewer。

    你现在应该看到一个执行的时间线。你可以使用 WASD 键来导航跟踪，点击或拖动以选择事件以查看底部的更多详细信息。有关使用跟踪查看器的更多详细信息，请参阅 这些 TensorFlow 文档。

    你也可以使用 memory_viewer、op_profile 和 graph_viewer 工具。

添加自定义跟踪事件#

默认情况下，跟踪查看器中的事件主要是低级别的内部 JAX 函数。您可以通过在代码中使用 jax.profiler.TraceAnnotation 和 jax.profiler.annotate_function() 来添加自己的事件和函数。

故障排除#

GPU 分析#

在GPU上运行的程序应在跟踪查看器的顶部附近生成GPU流的跟踪。如果你只看到主机跟踪，请检查你的程序日志和/或输出中是否有以下错误消息。

如果你遇到类似错误：无法加载动态库 'libcupti.so.10.1'
完整错误：

W external/org_tensorflow/tensorflow/stream_executor/platform/default/dso_loader.cc:55] Could not load dynamic library 'libcupti.so.10.1'; dlerror: libcupti.so.10.1: cannot open shared object file: No such file or directory
2020-06-12 13:19:59.822799: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1422] function cupti_interface_->Subscribe( &subscriber_, (CUpti_CallbackFunc)ApiCallback, this)failed with error CUPTI could not be loaded or symbol could not be found.

将 libcupti.so 的路径添加到环境变量 LD_LIBRARY_PATH 中。（尝试 locate libcupti.so 以找到路径。）例如：

export LD_LIBRARY_PATH=/usr/local/cuda-10.1/extras/CUPTI/lib64/:$LD_LIBRARY_PATH

如果在执行此操作后仍然收到 无法加载动态库 的消息，请检查GPU跟踪是否无论如何都显示在跟踪查看器中。即使一切正常，有时也会出现此消息，因为它会在多个位置查找 libcupti 库。

如果你遇到类似错误：failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
完整错误：

E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1445] function cupti_interface_->EnableCallback( 0 , subscriber_, CUPTI_CB_DOMAIN_DRIVER_API, cbid)failed with error CUPTI_ERROR_INSUFFICIENT_PRIVILEGES
2020-06-12 14:31:54.097791: E external/org_tensorflow/tensorflow/core/profiler/internal/gpu/cupti_tracer.cc:1487] function cupti_interface_->ActivityDisable(activity)failed with error CUPTI_ERROR_NOT_INITIALIZED

运行以下命令（注意这需要重启）：

echo 'options nvidia "NVreg_RestrictProfilingToAdminUsers=0"' | sudo tee -a /etc/modprobe.d/nvidia-kernel-common.conf
sudo update-initramfs -u
sudo reboot now

更多信息请参见 NVIDIA 关于此错误的文档。

在远程机器上进行性能分析#

如果你想要分析的 JAX 程序运行在远程机器上，一种选择是在远程机器上运行上述所有指令（特别是，在远程机器上启动 TensorBoard 服务器），然后使用 SSH 本地端口转发从本地机器访问 TensorBoard 的 Web UI。使用以下 SSH 命令将默认的 TensorBoard 端口 6006 从本地转发到远程机器：

ssh -L 6006:localhost:6006 <remote server address>

或者，如果您使用的是 Google Cloud：

$ gcloud compute ssh <machine-name> -- -L 6006:localhost:6006

多个 TensorBoard 安装#

如果启动 TensorBoard 时遇到类似错误：ValueError: Duplicate plugins for name projector

这通常是因为安装了两个版本的 TensorBoard 和/或 TensorFlow（例如，tensorflow、tf-nightly、tensorboard 和 tb-nightly pip 包都包含 TensorBoard）。卸载单个 pip 包可能会导致 tensorboard 可执行文件被移除，这很难替换，因此可能需要卸载所有内容并重新安装单个版本：

pip uninstall tensorflow tf-nightly tensorboard tb-nightly
pip install tensorflow

Nsight#

NVIDIA 的 Nsight 工具可以用于在 GPU 上追踪和分析 JAX 代码。详情请参阅 Nsight 文档。

目录