Mastering TensorFlow 1.x
上QQ阅读APP看书,第一时间看更新

GPU memory handling

When you start running the TensorFlow session, by default it grabs all of the GPU memory, even if you place the operations and variables only on one GPU in a multi-GPU system.  If you try to run another session at the same time,  you will get 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 executed after this setting will be able to grab all of the memory of only the visible GPU.

  • When you do not want the session to grab all of the memory of the GPU, then you can use the config option per_process_gpu_memory_fraction to allocate a percentage of memory:
config.gpu_options.per_process_gpu_memory_fraction = 0.5

This will allocate 50% of the memory of all the GPU devices.

  • You can also combine both of the above strategies, i.e. make only a percentage along with making only some of the GPU visible to the process.
  • You can also limit the TensorFlow process to grab only the minimum required memory at the start of the process. As the process executes further, you can set a config option to allow the growth of this memory.
config.gpu_options.allow_growth = True

This option only allows for the allocated memory to grow, but the memory is never released back.

You will learn techniques for distributing computation across multiple compute devices and multiple nodes in later chapters.