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

重置所选设备的跟踪内存统计信息。

用法

tf.config.experimental.reset_memory_stats(
    device
)

此函数将设备的跟踪峰值内存设置为设备的当前内存使用情况。这允许您测量程序特定部分的峰值内存使用量。例如：

if tf.config.list_physical_devices('GPU'):
  # Sets the peak memory to the current memory.
  tf.config.experimental.reset_memory_stats('GPU:0')
  # Creates the first peak memory usage.
  x1 = tf.ones(1000 * 1000, dtype=tf.float64)
  del x1 # Frees the memory referenced by `x1`.
  peak1 = tf.config.experimental.get_memory_info('GPU:0')['peak']
  # Sets the peak memory to the current memory again.
  tf.config.experimental.reset_memory_stats('GPU:0')
  # Creates the second peak memory usage.
  x2 = tf.ones(1000 * 1000, dtype=tf.float32)
  del x2
  peak2 = tf.config.experimental.get_memory_info('GPU:0')['peak']
  assert peak2 < peak1  # tf.float32 consumes less memory than tf.float64.

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