在一套標準的系統上通常有多個計算設備. TensorFlow 支持 CPU 和 GPU 這兩種設備. 我們用指定字符串 strings 來標識這些設備. 比如:
“/cpu:0”: 機器中的 CPU
“/gpu:0”: 機器中的 GPU, 如果你有一個的話.
“/gpu:1”: 機器中的第二個 GPU, 以此類推…
如果一個 TensorFlow 的 operation 中兼有 CPU 和 GPU 的實現, 當這個算子被指派設備時, GPU 有優先權. 比如matmul中 CPU 和 GPU kernel 函數都存在. 那麼在 cpu:0 和 gpu:0 中, matmul operation 會被指派給 gpu:0 .
# 新建一個graph.
with tf.device(’/cpu:0’):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name=‘a’)
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name=‘b’)
c = tf.matmul(a, b)
# 新建session with log_device_placement並設置爲True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 運行這個op.
print sess.run©
在多GPU系統裏使用單一GPU
with tf.device(’/gpu:2’):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name=‘a’)
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name=‘b’)
c = tf.matmul(a, b)
# 新建 session with log_device_placement 並設置爲 True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 運行這個 op.
print sess.run©
使用多個 GPU
c = []
for d in [’/gpu:2’, ‘/gpu:3’]:
with tf.device(d):
a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3])
b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2])
c.append(tf.matmul(a, b))
with tf.device(’/cpu:0’):
sum = tf.add_n©
# 新建session with log_device_placement並設置爲True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 運行這個op.
print sess.run(sum)