How to prevent tensorflow from allocating the totality of a GPU memory?

Я работаю в среде, в которой вычислительные ресурсы являются общими, т. Е. У нас есть несколько серверных машин, оснащенных несколькими графическими процессорами Nvidia Titan X каждая.

Для моделей малого и среднего размера 12 ГБ Titan X обычно достаточно для одновременного обучения 2-3 человек на одном графическом процессоре. Если модели достаточно малы, чтобы одна модель не использовала в полной мере все вычислительные блоки графического процессора, это может фактически привести к ускорению по сравнению с выполнением одного процесса обучения за другим. Даже в случаях, когда параллельный доступ к графическому процессору действительно замедляет время индивидуального обучения, все равно приятно иметь гибкость, заключающуюся в одновременном обучении нескольких пользователей на графическом процессоре.

Проблема с TensorFlow заключается в том, что по умолчанию он выделяет весь объем доступной памяти графического процессора при запуске. Даже для небольшой двухслойной нейронной сети я вижу, что все 12 ГБ памяти графического процессора израсходованы.

Есть ли способ заставить TensorFlow выделять только, скажем, 4 ГБ памяти графического процессора, если известно, что этого достаточно для данной модели?

Вы можете установить долю памяти графического процессора, которая будет выделяться при создании tf.Session путем передачи tf.GPUOptions как части необязательного config аргумента:

# Assume that you have 12GB of GPU memory and want to allocate ~4GB:
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)

sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))

per_process_gpu_memory_fraction действует как жесткая верхняя граница объема памяти графического процессора, который будет использоваться процессом на каждом графическом процессоре на одной машине. В настоящее время эта доля применяется равномерно ко всем графическим процессорам на одной машине; нет способа установить это для каждого графического процессора.

config = tf.ConfigProto()
config.gpu_options.allow_growth=True
sess = tf.Session(config=config)

https://github.com/tensorflow/tensorflow/issues/1578

Для TensorFlow 2.0 и 2.1 (документы):

import tensorflow as tf
tf.config.gpu.set_per_process_memory_growth(True)

Для TensorFlow 2.2+ (документы):

import tensorflow as tf
gpus = tf.config.experimental.list_physical_devices('GPU')
for gpu in gpus:
  tf.config.experimental.set_memory_growth(gpu, True)

В документах также перечислены еще несколько методов:

• Установите переменную окружения TF_FORCE_GPU_ALLOW_GROWTH в true.


• Используйте tf.config.experimental.set_virtual_device_configuration для установки жесткого ограничения на виртуальное устройство с графическим процессором.

Вот выдержка из книги Deep Learning with TensorFlow

В некоторых случаях желательно, чтобы процесс выделял только подмножество доступной памяти или увеличивал использование памяти только по мере необходимости процесса. TensorFlow предоставляет два варианта конфигурации сеанса для управления этим. Первый - это allow_growth вариант, который пытается выделить только столько памяти графического процессора на основе распределения во время выполнения, вначале выделяется очень мало памяти, и по мере запуска сеансов и увеличения объема памяти графического процессора мы расширяем область памяти графического процессора, необходимую процессу TensorFlow.

1) Разрешить рост: (более гибкий)

config = tf.ConfigProto()
config.gpu_options.allow_growth = True
session = tf.Session(config=config, ...)

Второй метод - это per_process_gpu_memory_fraction option, который определяет долю от общего объема памяти, которая должна быть выделена each видимому графическому процессору. Примечание: освобождение памяти не требуется, это может даже ухудшить фрагментацию памяти, когда это будет сделано.

2) Выделить фиксированную память:

Выделять только 40% от общей памяти каждого графического процессора с помощью:

config = tf.ConfigProto()
config.gpu_options.per_process_gpu_memory_fraction = 0.4
session = tf.Session(config=config, ...)

Примечание:
Это полезно только в том случае, если вы действительно хотите привязать объем памяти графического процессора, доступный процессу TensorFlow.