jax.profiler.start_trace#

jax.profiler.start_trace(log_dir, create_perfetto_link=False, create_perfetto_trace=False, profiler_options=None)[源代码]#

启动一个性能分析器跟踪。

此跟踪将捕获 CPU、GPU 和/或 TPU 活动，包括 Python 函数和 JAX 设备上的操作。使用 stop_trace() 来结束跟踪并将结果保存到 log_dir。

生成的跟踪可以使用 TensorBoard 查看。请注意，在收集跟踪时不需要运行 TensorBoard。

一次只能收集一个跟踪。如果当前有一个跟踪正在运行而调用 start_trace()，则会引发 RuntimeError。

参数:

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

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

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

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

返回类型:

无