tf.device()指定tensorflow運行的GPU或CPU設備

在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 

發表評論
所有評論
還沒有人評論,想成為第一個評論的人麼? 請在上方評論欄輸入並且點擊發布.
相關文章