在tensorflow中，我們可以使用 tf.device() 指定模型運(yùn)行的具體設(shè)備，可以指定運(yùn)行在GPU還是CUP上，以及哪塊GPU上。
設(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()指定tensorflow運(yùn)行的GPU或CPU設(shè)備實(shí)現(xiàn)的文章就介紹到這了,更多相關(guān)tensorflow運(yùn)行GPU或CPU內(nèi)容請(qǐng)搜索腳本之家以前的文章或繼續(xù)瀏覽下面的相關(guān)文章希望大家以后多多支持腳本之家！
您可能感興趣的文章:- tensorflow 利用expand_dims和squeeze擴(kuò)展和壓縮tensor維度方式
- TensorFlow低版本代碼自動(dòng)升級(jí)為1.0版本
- TensorFlow的環(huán)境配置與安裝方法
- TensorFlow2.0使用keras訓(xùn)練模型的實(shí)現(xiàn)
- 解決tensorflow模型壓縮的問題_踩坑無(wú)數(shù),總算搞定
