详解tf.device()指定tensorflow运行的GPU或CPU设备实现

时间:2021-05-22

在tensorflow中,我们可以使用tf.device()指定模型运行的具体设备,可以指定运行在GPU还是CUP上,以及哪块GPU上。

设置使用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)

到此这篇关于详解tf.device()指定tensorflow运行的GPU或CPU设备实现的文章就介绍到这了,更多相关tensorflow运行GPU或CPU内容请搜索以前的文章或继续浏览下面的相关文章希望大家以后多多支持!

声明:本页内容来源网络,仅供用户参考;我单位不保证亦不表示资料全面及准确无误,也不保证亦不表示这些资料为最新信息,如因任何原因,本网内容或者用户因倚赖本网内容造成任何损失或损害,我单位将不会负任何法律责任。如涉及版权问题,请提交至online#300.cn邮箱联系删除。

相关文章