jax.profiler.保存设备内存配置文件

jax.profiler.保存设备内存配置文件#

jax.profiler.save_device_memory_profile(filename, backend=None)[源代码][源代码]#

收集设备内存配置文件并将其写入文件。

save_device_memory_profile() 是围绕 device_memory_profile() 的一个便捷包装，它将其输出保存到 filename 中。更多信息请参阅 device_memory_profile() 文档。

参数:

  • filename – 应写入配置文件的文件名。

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

返回类型:

None

目录