使用tf.device()怎么指定運(yùn)行tensorflow

這篇文章給大家介紹使用tf.device()怎么指定運(yùn)行tensorflow，內(nèi)容非常詳細(xì)，感興趣的小伙伴們可以參考借鑒，希望對(duì)大家能有所幫助。

設(shè)置使用GPU

使用 tf.device('/gpu:1') 指定Session在第二塊GPU上運(yùn)行：

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() 中參數(shù) log_device_placement=True  會(huì)打印出執(zhí)行操作所用的設(shè)備，以上輸出：

如果安裝的是GPU版本的tensorflow，機(jī)器上有支持的GPU，也正確安裝了顯卡驅(qū)動(dòng)、CUDA和cuDNN，默認(rèn)情況下，Session會(huì)在GPU上運(yùn)行：

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)

默認(rèn)在GPU:0上執(zhí)行：

設(shè)置使用cpu

tensorflow中不同的GPU使用/gpu:0和/gpu:1區(qū)分，而CPU不區(qū)分設(shè)備號(hào)，統(tǒng)一使用 /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)

關(guān)于使用tf.device()怎么指定運(yùn)行tensorflow就分享到這里了，希望以上內(nèi)容可以對(duì)大家有一定的幫助，可以學(xué)到更多知識(shí)。如果覺得文章不錯(cuò)，可以把它分享出去讓更多的人看到。