jax.profiler.start_trace

内容

jax.profiler.start_trace#

jax.profiler.start_trace(log_dir, create_perfetto_link=False, create_perfetto_trace=False)[source]#

启动性能分析跟踪。

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

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

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

参数：

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

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

  • create_perfetto_trace (bool) – 如果为真，则此布尔值还会转储一个与 Perfetto 追踪查看器 UI (https://ui.perfetto.dev) 兼容的 perfetto_trace.json.gz 文件。如果 create_perfetto_link 为真，则也会生成该文件。如果您想生成一个兼容 Perfetto 的追踪而无需阻塞进程，这可能很有用。

返回类型:

无

内容