首先来描述一下我遇到的问题。最近想做一个训练实验,比较两个不同超参数配置下的训练效果。我比较了一下使用GPU和使用CPU的训练速度,差距还是很明显的。但是,问题是我只有一台工作站,配有一块Quadro K2200d的显卡,显存是4G,如果我启动一个训练进程,显存占用空间就可以到达3G,但是GPU的占用率却很低,大部分时间都是0。
查看GPU使用状况
nvidia-smi -l --loop=2
使用以上命令可以持续跟踪GPU的使用情况,每2秒更新一次。
如果可以解决显存不足的问题,让两个训练都使用GPU的话,就可以提高GPU的使用效率,也提高整个训练的效率。
解决这个问题的方法很简单, 只需要在tf.Session里设定该训练使用GPU Memory的比例,让两个训练同时运行时显存不会不够即可。设置GPU Memory使用比例的方法如下。
分配显存给不同训练进程
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.35)
with tf.Session(config=tf.ConfigProto(log_device_placement=True, gpu_options=gpu_options)) as sess:
#你的训练代码
我把两个训练的GPU Memory占用比例都设置成35%, 两个训练就都可以跑在GPU上了。此处设置占比时还需要考虑其他进程也需要占用GPU,所以不可以设置为50%。我还尝试过设置成40%, 但因为gnome-shell, Xorg等进程占用了将近300M的显存,我实际可用的显存大约只有3.3G左右。
分配GPU给不同训练进程
如果你有多块GPU的话,也可以在指定训练进程使用的GPU, 在导入tensorflow之前 指定可见的GPU,比如以下代码则指定改训练仅可以使用GPU 0 和1。
import os
os.environ["CUDA_VISIBLE_DEVICES"]="0","1"
import tensorflow as tf
如果不想使用GPU,则把GPU资源指定为-1即可,代码如下
import os
os.environ["CUDA_VISIBLE_DEVICES"]="-1"
import tensorflow as tf