远程性能分析#

当分析远程运行的代码（例如在托管 VM 上）时，您需要为链接工作在端口 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=<directory of choice> 来覆盖。此外，默认情况下，程序将提示您打开指向 ui.perfetto.dev 的链接。当您打开链接时，Perfetto UI 将加载跟踪文件并打开可视化工具。通过将 --no_perfetto_link 传入命令，可以禁用此功能。或者，您也可以将 Tensorboard 指向 log_dir 以分析跟踪（请参阅下面的“Tensorboard 性能分析”部分）。

安装#

TensorBoard 性能分析器仅适用于与 TensorFlow 捆绑的 TensorBoard 版本。

pip install tensorflow tensorboard-plugin-profile

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

TensorBoard 性能分析器的 nightly 版本需要 nightly tensorflow 和 tensorboard

pip install tf-nightly tb-nightly tbp-nightly

程序化捕获#

您可以使用 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”，或直接转到 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 挂钩，因此最初可能看起来像没有捕获任何内容。）在左侧的“Tools”下，选择 trace_viewer。

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

   您还可以使用 memory_viewer、op_profile 和 graph_viewer 工具。
   
   

添加自定义跟踪事件#

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

故障排除#

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.

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

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

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

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

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