轉自:牧野
設置使用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)