TensorFlow指定GPU()

1、指定某个显卡
可以在文件开头加入如下代码:

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1, 2"  # 使用第二,三块GPU(从0开始)

禁用GPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

支持的设备表示:

"/cpu:0":机器的 CPU。
"/device:GPU:0":机器的 GPU(如果有一个)。
"/device:GPU:1":机器的第二个 GPU(以此类推)。

2、手动分配设备

# Creates a graph.
with tf.device('/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)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print(sess.run(c))

3、使用多个 GPU

如果您想要在多个 GPU 上运行 TensorFlow,则可以采用多塔式方式构建模型,其中每个塔都会分配给不同 GPU。例如:

# Creates a graph.
c = []
for d in ['/device:GPU:2', '/device: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)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print(sess.run(sum))
————————

1、指定某个显卡
可以在文件开头加入如下代码:

import os
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = "1, 2"  # 使用第二,三块GPU(从0开始)

禁用GPU

import os
os.environ["CUDA_VISIBLE_DEVICES"] = "-1"

支持的设备表示:

"/cpu:0":机器的 CPU。
"/device:GPU:0":机器的 GPU(如果有一个)。
"/device:GPU:1":机器的第二个 GPU(以此类推)。

2、手动分配设备

# Creates a graph.
with tf.device('/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)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print(sess.run(c))

3、使用多个 GPU

如果您想要在多个 GPU 上运行 TensorFlow,则可以采用多塔式方式构建模型,其中每个塔都会分配给不同 GPU。例如:

# Creates a graph.
c = []
for d in ['/device:GPU:2', '/device: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)
# Creates a session with log_device_placement set to True.
sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))
# Runs the op.
print(sess.run(sum))