Skip to main content
 首页 » 编程设计

python之我如何使用 2 gpu 在 tensorflow 中进行计算

2025年02月15日24shanyou

我使用两个 GTX 980 GPU。 当我在 tensorflow 中处理 slim 时。 通常,我遇到所谓的“内存不足”问题。 所以,我想同时使用两个 gpu。 我怎样才能使用2个GPU? 哦,对不起,我的英语水平很差。 :(

请您参考如下方法:

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

 # Creates a 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) 
# 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)) 

您将看到以下输出。

    Device mapping: 
/job:localhost/replica:0/task:0/gpu:0 -> device: 0, name: Tesla K20m, pci bus 
id: 0000:02:00.0 
/job:localhost/replica:0/task:0/gpu:1 -> device: 1, name: Tesla K20m, pci bus 
id: 0000:03:00.0 
/job:localhost/replica:0/task:0/gpu:2 -> device: 2, name: Tesla K20m, pci bus 
id: 0000:83:00.0 
/job:localhost/replica:0/task:0/gpu:3 -> device: 3, name: Tesla K20m, pci bus 
id: 0000:84:00.0 
Const_3: /job:localhost/replica:0/task:0/gpu:3 
Const_2: /job:localhost/replica:0/task:0/gpu:3 
MatMul_1: /job:localhost/replica:0/task:0/gpu:3 
Const_1: /job:localhost/replica:0/task:0/gpu:2 
Const: /job:localhost/replica:0/task:0/gpu:2 
MatMul: /job:localhost/replica:0/task:0/gpu:2 
AddN: /job:localhost/replica:0/task:0/cpu:0 
[[  44.   56.] 
 [  98.  128.]] 

我不确定“内存不足”问题,因为您没有提供代码片段。