jax.profiler.trace#

jax.profiler.trace(log_dir, create_perfetto_link=False, create_perfetto_trace=False, profiler_options=None)[source]#

用于生成性能分析器跟踪的上下文管理器。

此跟踪将捕获 CPU、GPU 和/或 TPU 活动，包括 Python 函数和 JAX 设备上的操作。

生成的跟踪可以在 TensorBoard 中查看。请注意，在收集跟踪时，TensorBoard 无需运行。

一次只能收集一个跟踪。如果在另一个跟踪正在运行时启动跟踪，则会引发 RuntimeError。

参数:

  • log_dir (os.PathLike | str) – 保存性能分析器跟踪的目录（通常是 TensorBoard 日志目录）。

  • create_perfetto_link – 一个布尔值，如果为 True，则创建并打印指向 Perfetto 跟踪查看器 UI (https://ui.perfetto.dev) 的链接。程序将阻塞，直到链接被打开且 Perfetto 加载跟踪。

  • create_perfetto_trace – 一个布尔值，如果为 True，则额外转储一个 perfetto_trace.json.gz 文件，该文件兼容上传到 Perfetto 跟踪查看器 UI (https://ui.perfetto.dev)。如果 create_perfetto_link 为 True，也会生成该文件。如果您想生成一个兼容 Perfetto 的跟踪而不阻塞进程，这可能会很有用。

  • profiler_options (ProfileOptions | None) – 用于配置性能分析器以进行收集的选项。