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 設備上調用，則會引發異常。