Как предотвратить tensorflow от выделения всей памяти GPU?

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

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

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

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

5 169
python tensorflow nvidia-titan

5 ответов:

вы можете установить долю памяти GPU, которая будет выделена при создании tf.Session, передав tf.GPUOptions на

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

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

123

вот отрывок из книги Deep Learning with TensorFlow

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

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

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

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

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

выделить только 40% из общей памяти каждого GPU по:

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

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

17

все ответы выше предполагают выполнение с sess.run() вызов, который становится исключением, а не правилом в последних версиях TensorFlow.

при использовании tf.Estimator framework (TensorFlow 1.4 и выше) способ передачи фракции вдоль неявно созданного MonitoredTrainingSession есть

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
trainingConfig = tf.estimator.RunConfig(session_config=conf, ...)
tf.estimator.Estimator(model_fn=..., 
                       config=trainingConfig)

аналогично в режиме ожидания (TensorFlow 1.5 и выше),

opts = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)
conf = tf.ConfigProto(gpu_options=opts)
tfe.enable_eager_execution(config=conf)

Edit: 11-04-2018 Например, если вы используете tf.contrib.gan.train, тогда вы можете использовать что-то похожее на ниже:

tf.contrib.gan.gan_train(........, config=conf)
8

бесстыдный плагин: если вы установите поддерживаемый GPU Tensorflow, сеанс сначала выделит все GPU, независимо от того, используете ли вы его только CPU или GPU. Я могу добавить свой совет, что даже если вы установите график для использования только CPU, вы должны установить ту же конфигурацию(как ответили выше:)), чтобы предотвратить нежелательное занятие GPU.

и в интерактивном интерфейсе, таком как IPython, вы также должны установить эту настройку, иначе она выделит всю память и не оставит почти ничего для других. Это иногда это трудно заметить.

3