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

You can set the fraction of GPU memory to be allocated when you construct a tf.Session by passing a tf.GPUOptions as part of the optional config argument:

# 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))

The per_process_gpu_memory_fraction acts as a hard upper bound on the amount of GPU memory that will be used by the process on each GPU on the same machine. Currently, this fraction is applied uniformly to all of the GPUs on the same machine; there is no way to set this on a per-GPU basis.


For TensorFlow 2.0 and 2.1 (docs):

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

For TensorFlow 2.2+ (docs):

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

The docs also list some more methods:

  • Set environment variable TF_FORCE_GPU_ALLOW_GROWTH to true.
  • Use tf.config.experimental.set_virtual_device_configuration to set a hard limit on a Virtual GPU device.

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

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