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

Posted on by

在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內容請搜索WalkonNet以前的文章或繼續瀏覽下面的相關文章希望大傢以後多多支持WalkonNet！

推薦閱讀：