【TensorFlow】TensorFlow指定特定GPU以及占用显存的比例

因为LZ是使用GPU服务器跑TensorFlow，而TensorFlow默认的是占用所有GPU，于是为了不影响其他同学使用GPU，于是就试验和总结了一下TensorFlow指定GPU的方法。。

环境
系统：Ubuntu14.04

TensorFlow：v1.3

GPU 8个GTX1080，第一列的0~7的数是GPU的序号

 

一.设置指定GPU
1.一劳永逸的方法，直接在~/.bashrc中设置环境变量
CUDA_VISIBLE_DEVICES

      （1）vim ~/.bashrc

      （2）在.bashrc.最下方添加

export  CUDA_VISIBLE_DEVICES='XX'(这个xx是GPU的序号，比如我这边服务器上有8个GTX1080序号是0~7，可以填入任意多个序号，序号间用逗号隔开)

      （3）source ~/.bashrc

         这样只要每次你用TensorFlow跑程序，都只会占用指定的GPU

2.在终端设置使用的GPU
如果用方法一，虽然方便，但有的时候还是需要指定其他的GPU，这时可以这样，例如

CUDA_VISIBLE_DEVICES=2 python test.py
这样就只会使用序号为2的GPU

3.在程序中指定使用的GPU
import os

os.environ["CUDA_VISIBLE_DEVICES"]=‘6‘’，‘7’

这里就设置了使用序号为6,7两个的GPU

事实上这3种方法都是为了设置环境变量CUDA_VISIBLE_DEVICES，方法一可以分别和方法二，三结合使用，比如虽然方法一设置了CUDA_VISIBLE_DEVICES=2，后面再用方法二设置使用GPU 0，执行的时候只使用GPU 0。

二.设置占用每个GPU的比例

1.定量设置显存
gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=0.7)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options)) 
      这样运行TensorFlow程序时，每个使用的GPU中，占用的显存都不超过总显存的0.7.
      比如我的GTX1080显存是8G左右，然后占用的显存应该是0.7*8=5.6G左右，得到的结果5.9G，差不多。

2.按需设置显存
gpu_options = tf.GPUOptions(allow_growth=True)
sess = tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))   
    这样设置以后，程序就会按需占用GPU显存。