tf.device()指定运行设备
在TensorFlow中,模型可以在本地的GPU和CPU中运行,用户可以指定模型运行的设备。通常,如果你的TensorFlow版本是GPU版本的,而且你的电脑上配置有符合条件的显卡,那么在不做任何配置的情况下,模型是默认运行在显卡下的。
如果需要切换成CPU运算,可以调用tf.device(device_name)
函数,其中device_name
格式如/cpu:0
其中的0
表示设备号,TF不区分CPU的设备号,设置为0
即可。GPU区分设备号/gpu:0
和/gpu:1
表示两张不同的显卡。
在一些情况下,我们即使是在GPU下跑模型,也会将部分Tensor储存在内存里,因为这个Tensor可能太大了,显存不够放,相比于显存,内存一般大多了,于是这个时候就常常人为指定为CPU设备。这种形式我们在一些代码中能见到。如:
with tf.device('/cpu:0'):
build_CNN() # 此时,这个CNN的Tensor是储存在内存里的,而非显存里。
需要注意的是,这个方法会减少显存的负担,但是从内存把数据传输到显存中是非常慢的,这样做常常会减慢速度。
设置使用GPU
使用 tf.device(’/gpu:1’) 指定Session在第二块GPU上运行:
import tensorflow as tf
with tf.device('/gpu:1'):
v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v1')
v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v2')
sumV12 = v1 + v2
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
print sess.run(sumV12)
ConfigProto() 中参数 log_device_placement=True 会打印出执行操作所用的设备,以上输出:
如果安装的是GPU版本的tensorflow,机器上有支持的GPU,也正确安装了显卡驱动、CUDA和cuDNN,默认情况下,Session会在GPU上运行:
import tensorflow as tf
v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v1')
v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v2')
sumV12 = v1 + v2
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
print sess.run(sumV12)
默认在GPU:0上执行:
设置使用cpu
tensorflow中不同的GPU使用/gpu:0和/gpu:1区分,而CPU不区分设备号,统一使用 /cpu:0
import tensorflow as tf
with tf.device('/cpu:0'):
v1 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v1')
v2 = tf.constant([1.0, 2.0, 3.0], shape=[3], name='v2')
sumV12 = v1 + v2
with tf.Session(config=tf.ConfigProto(log_device_placement=True)) as sess:
print sess.run(sumV12)
参考:https://blog.youkuaiyun.com/dcrmg/article/details/79747882