設定使用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

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)