Функция config.experimental.set_memory_growth

Функция config.experimental.set_memory_growth управляет режимом динамического роста памяти для указанного физического устройства GPU. По умолчанию TensorFlow выделяет всю доступную память GPU при первой инициализации, что может привести к ошибкам нехватки памяти при запуске нескольких процессов. Данная функция позволяет выделять память по мере необходимости. Первым параметром передается объект физического устройства, полученный через config.list_physical_devices, вторым - логическое значение True или False, включающее или отключающее режим роста.

Синтаксис

tf.config.experimental.set_memory_growth(device, enable)

Пример

Давайте включим режим динамического роста памяти для первого доступного GPU:

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
if gpus:
	tf.config.experimental.set_memory_growth(gpus[0], True)
	print("Memory growth enabled for", gpus[0].name)

Результат выполнения кода:

"Memory growth enabled for /physical_device:GPU:0"

Пример

Давайте проверим текущий статус режима роста памяти с помощью парной функции config.experimental.get_memory_growth:

<+python+> import tensorflow as tf gpus = tf.config.list_physical_devices('GPU') if gpus: tf.config.experimental.set_memory_growth(gpus[0], True) status = tf.config.experimental.get_memory_growth(gpus[0]) print("Memory growth status:", status) <-python+>

Результат выполнения кода:

"Memory growth status: True"

Пример

Давайте включим режим роста памяти для всех доступных GPU в цикле:

import tensorflow as tf

gpus = tf.config.list_physical_devices('GPU')
for gpu in gpus:
	tf.config.experimental.set_memory_growth(gpu, True)
	print("Configured:", gpu.name)

Результат выполнения кода:

"Configured: /physical_device:GPU:0"

Смотрите также

функцию get_memory_growth,
которая возвращает текущий статус режима роста памяти
функцию list_physical_devices,
которая возвращает список доступных физических устройств
функцию set_visible_devices,
которая ограничивает набор видимых устройств
функцию set_virtual_device_configuration,
которая настраивает виртуальные устройства