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.

В будущем могут быть добавлены дополнительные ключи, включая ключи, специфичные для устройства.

В настоящее время поддерживаются только устройства GPU и TPU. Если функция вызвана для устройства CPU, будет возбуждено исключение.

Для GPU TensorFlow по умолчанию выделяет всю память, если не указано иное с помощью tf.config.experimental.set_memory_growth. Словарь содержит только текущее и максимальное использование памяти, которое фактически используется TensorFlow, а не всю выделенную память на GPU.