jax.profiler 模块#

追踪和时间性能分析#

性能分析计算 描述了如何使用 JAX 的追踪和时间性能分析功能。

start_server(端口)

在端口 port 上启动 profiler 服务器。

start_trace(log_dir[, create_perfetto_link, ...])

启动 profiler 追踪。

stop_trace()

停止当前正在运行的 profiler 追踪。

trace(log_dir[, create_perfetto_link, ...])

用于进行 profiler 追踪的上下文管理器。

annotate_function(func[, name])

为函数执行生成追踪事件的装饰器。

TraceAnnotation

在 profiler 中生成追踪事件的上下文管理器。

StepTraceAnnotation(name, **kwargs)

在 profiler 中生成步骤追踪事件的上下文管理器。

设备内存性能分析#

请参阅 性能分析设备内存 以了解 JAX 设备内存性能分析功能的介绍。

device_memory_profile([backend])

将 JAX 设备内存性能分析捕获为 pprof 格式的协议缓冲区。

save_device_memory_profile(filename[, backend])

收集设备内存性能分析并将其写入文件。