TensorFlow Machine Learning Projects

GPU memory handling

At the start of the TensorFlow session, by default, a session grabs all of the GPU memory, even if the operations and variables are placed only on one GPU in a multi-GPU system. If another session starts execution at the same time, it will receive an out-of-memory error. This can be solved in multiple ways:

• For multi-GPU systems, set the environment variable CUDA_VISIBLE_DEVICES=<list of device idx>:

os.environ['CUDA_VISIBLE_DEVICES']='0'

The code that's executed after this setting will be able to grab all of the memory of the visible GPU.

• For letting the session grab a part of the memory of the GPU, use the config option per_process_gpu_memory_fraction to allocate a percentage of the memory:

config.gpu_options.per_process_gpu_memory_fraction ...