jax.profiler.设备内存概况

jax.profiler.设备内存概况#

jax.profiler.device_memory_profile(backend=None)[源代码][源代码]#

捕获 JAX 设备内存配置文件为 pprof 格式的协议缓冲区。

设备内存概况是内存状态的快照，描述了内存中存在的 JAX Array 和可执行对象及其分配位置。

有关如何使用设备内存分析器的更多信息，请参阅 分析设备内存。

分析系统通过检测 JAX 设备上的分配来工作，为每次分配捕获 Python 堆栈跟踪。检测始终启用；device_memory_profile() 提供了一个 API 来捕获它。

The output of device_memory_profile() is a binary protocol buffer that can be interpreted and visualized by the pprof tool.

参数:

backend (str | None) – 可选；应收集设备内存配置文件的 JAX 后端名称。

返回:

包含二进制 pprof 格式协议缓冲区的字节字符串。

返回类型:

bytes

目录