一、使用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)