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

32 篇文章 0 订阅
20 篇文章 1 订阅

在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)

 

---------------------  
作者:-牧野-  
来源:CSDN  
原文:https://blog.csdn.net/dcrmg/article/details/79747882 

  • 1
    点赞
  • 3
    收藏
    觉得还不错? 一键收藏
  • 0
    评论
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值