在tensorflow中,我们可以使用 tf.device() 指定模型运行的具体设备,可以指定运行在GPU还是CUP上,以及哪块GPU上。
![](https://i-blog.csdnimg.cn/blog_migrate/810bd31a5da8c1db7da3a96c5fdb5df7.jpeg)
如果安装的是GPU版本的tensorflow,机器上有支持的GPU,也正确安装了显卡驱动、CUDA和cuDNN,默认情况下,Session会在GPU上运行:
![](https://i-blog.csdnimg.cn/blog_migrate/522eea2ef664cda427cd1c5a8ffc3bf8.jpeg)
设置使用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 会打印出执行操作所用的设备,以上输出:
![](https://i-blog.csdnimg.cn/blog_migrate/810bd31a5da8c1db7da3a96c5fdb5df7.jpeg)
如果安装的是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上执行:
![](https://i-blog.csdnimg.cn/blog_migrate/522eea2ef664cda427cd1c5a8ffc3bf8.jpeg)
设置使用cpu
tensorflow中不同的GPU使用/gpu:0和/gpu:1区分,而CPU不区分设备号,统一使用 /cpu:0import 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://i-blog.csdnimg.cn/blog_migrate/4c9a40c8267f52448c60ac38dcf5f044.jpeg)