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, будет возбуждено исключение.