使用tensorflow-gpu时设置GPU

版权声明：本文为博主原创文章，未经博主允许不得转载。 https://blog.csdn.net/weixin_34613450/article/details/82315277

使用GPU跑tensorflow程序，默认加载所有的GPU，但计算过程中只会用其中一块。也就是你看着所有GPU都被占用了，以为是在GPU并行计算，但实际上只有其中一块在运行；另外的所有显卡都闲着，但其显存都被占用了，所以别人也用不了。不过这种情况通过在程序之前加三行代码就可以解决：

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ['CUDA_VISIBLE_DEVICES'] = "0"

这行代码加在TensorFlow程序开头，就可以成功屏蔽掉系统中除 gpu0 （当然，这个gpu序号要根据实际情况来定）之外所有的GPU设备了。

如果想使用GPU0和GPU1，可以将第三行代码改为：

os.environ['CUDA_VISIBLE_DEVICES'] = "0,1"

注意，第二行 os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" 也很重要，保证程序中的GPU序号是和硬件中的序号是相同的，否则可能会造成不少的麻烦。

如果不想使用GPU，也有办法。这样设置第三行代码 os.environ['CUDA_VISIBLE_DEVICES'] = "" ，这样这个程序就不能看见所有的GPU了。

除此之外，TensorFlow程序会吃掉所用显卡的所有显存，如果想让程序需要多少显存就用多少应该怎么设置呢？创建 session 的时候加一项设置：

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

程序运行过程中查看显存使用情况：

nvidia-smi