【深度学习框架】TensorFlow 选择计算设备 CPU/GPU

在计算实现上, TensorFlow 将图形定义转换成分布式执行的操作, 以充分利用可用的计算资源。

TensorFlow 支持 CPU 和 GPU 两种计算设备。

一般你不需要显式指定使用 CPU 还是 GPU, TensorFlow 能自动检测。如果检测到 GPU, TensorFlow 会优先利用找到的第一个 GPU 来执行操作。如果机器上有超过一个可用的 GPU, 除第一个外的其它 GPU 默认是不参与计算的。

如果你希望手工指派设备，你可以用 with tf.device 创建一个设备环境, 这个环境下的 operation 都统一运行在环境指定的设备上。

设备用字符串进行标识：

• “/cpu:0”: 机器的 CPU.
• “/gpu:0”: 机器的第一个 GPU, 如果有的话.
• “/gpu:1”: 机器的第二个 GPU, 以此类推.

在多GPU系统里使用单一GPU

如果你的系统里有多个 GPU, 那么 ID 最小的 GPU 会默认使用. 如果你想用别的 GPU, 可以用下面的方法显式的声明你的偏好：

# 新建一个 graph.
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(c)

使用多个 GPU

如果你想让 TensorFlow 使用电脑上的GPU集群同时进行运算， 你需要将 op 明确指派 给不同的GPU执行。

# 新建一个 graph.
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(c)
# 新建session with log_device_placement并设置为True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# 运行这个op.
print sess.run(sum)

---

参考：TensorFlow中文社区 文档