Tensorflow默认占满全部GPU的全部资源

Posted 2021-01-10 weiyinfu

一台服务器上装了多块GPU，默认情况下启动一个深度学习训练任务时，这个任务会占满每一块GPU的几乎全部存储空间。这就导致一个服务器基本上只能执行一个任务，而实际上任务可能并不需要如此多的资源，这相当于一种资源浪费。

针对这个问题，有如下解决方案。

一、直接设置可见GPU

写一个脚本，设置环境变量

export CUDA_VISIBLE_DEVICES=0
python model.py

二、设置每个GPU的存储上限

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.333)

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

三、增量使用GPU存储

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