tensorflow 显存 训练_tensorflow手动指定GPU以及显存大小

以前我们组就一块显卡，不存在指定设备的问题。近期刚插了一块新的gtx 1080ti，几人公用两块卡来做训练、测试、预测等等，网上找了个方式可以指定使用的设备，并且限定使用的显存大小，还是很有用的，亲测可行，拿过来分享给大家~~~~~

为了测试方便，使用了mnist代码，数据集和代码都比较小，比较好控制。mnist代码有点多就不贴了，大家自己找找吧。

下面开始正题：

一般来说GPU device编号是从0开始的，第二块的话就是1。如果不指定，那tf就优先使用GPU并且默认用第一块卡。

可以通过下面的方式来手动指定使用的GPU：

import os

os.environ["CUDA_VISIBLE_DEVICES"] = "1"

使用nvidia-smi命令查看一下GPU的使用情况：

发现使用的设备已经是GPU1了，并且显存是用满的。

然后我们下面的代码来限定显存使用：

gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)

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

我们设置使用显存为实际显存的70%，看看效果如何：