jax.profiler.trace

jax.profiler.trace#

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

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

跟踪将捕获 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 跟踪查看器 UI (https://ui.perfetto.dev) 兼容的 `perfetto_trace.json.gz` 文件以供上传。如果 `create_perfetto_link` 为 true，也将生成该文件。如果您想生成一个与 Perfetto 兼容的跟踪，而无需阻塞进程，这可能很有用。