TensorFlow指定GPU设备、分配显存

一、使用GPU加速运算

        1.TensorFlow程序中可以通过tf.device函数来指定运行的设备。

          CPU名称为"/cpu:0"，即使有多个CPU，所有的CPU都使用/cpu:0作为名称。

          GPU则不同，是"/gpu:n"，n可替换。n=0，即是第一个GPU，以此类推。

        2.观察每个运算所用的设备：

          使用log_device_placement传入会话当中的配置参数。具体代码：

            

import tensorflow as tf
sess=tf.Session(config=tf.ConfigProto(log_devic_placement=True))
sess.run(...)    

          运行程序后会显示每一个操作所用的设备显示到屏幕上：

          

          如果没有明确指定运行设备，那么TensorFlow会优先选择GPU。

        2.指定运算到指定的设备

          需要tf.device来手工指定。代码如下：

          

          得到如下结果：

          

          注意：有些操作是不能在GPU上完成，否则会报错。

              比如：tf.Variable()操作只支持实数型数据的参数在GPU上运行，整数型就会报错。

              具体哪些不能在GPU上运行，参考：官方部分文档

        为了避免不知道某个操作是否确定在GPU上运行，可指定GPU设备后，在构建会话时，传入参数：                allow_soft_placement=True，使得不能在GPU设备上完成的自动移到CPU上完成，具体代码：

              

sess = tf.Session(config=tf.ConfigProto(allow_soft_placement=True))

二、TensorFlow动态分配GPU和显存

          TensorFlow会默认占用所有GPU以及每个GPU的所有显存。

          1.如果只想使用部分GPU，2种方法进行设置：

            第一种：

            运行.py程序时指定：CUDA_VISIBLE_DEVICES=0, 1 python xx.py

            这时候就会只有第一块、第二块GPU会使用。

            第二种：

            在程序中指定：

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "2"
#只使用第三块GPU

          2.动态分配显存

          让TensorFlow按需分配显存：      

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

          直接指定显存比例：

config.gpu_options.per_process_gpu_memory_fraction = 0.4
#占用GPU40%

           再传入会话当中：

session = tf.Session(config=config)