Python tf.config.experimental.get_memory_info用法及代码示例

获取所选设备的内存信息，作为字典。

用法

tf.config.experimental.get_memory_info(
    device
)

这个函数返回一个包含设备内存使用信息的字典。例如：

if tf.config.list_physical_devices('GPU'):
  # Returns a dict in the form {'current':<current mem usage>,
  #                             'peak':<peak mem usage>}
  tf.config.experimental.get_memory_info('GPU:0')

当前返回以下键：

• 'current' ：设备当前使用的内存，以字节为单位。
• 'peak' ：设备在程序运行期间使用的峰值内存，以字节为单位。可以使用 tf.config.experimental.reset_memory_stats 重置。

未来可能会添加更多键，包括device-specific 键。

目前仅支持 GPU 和 TPU。如果在 CPU 设备上调用，则会引发异常。

对于 GPU，TensorFlow 将默认分配所有内存，除非更改为 tf.config.experimental.set_memory_growth 。 dict 仅指定 TensorFlow 实际使用的当前和峰值内存，而不是 TensorFlow 在 GPU 上分配的内存。