7

我想在 2 个 GPU 上并行化以下简单表达式：在对结果求和之前在 GPU 0 和GPU 1 上进行C = A^n + B^n计算。A^nB^n

在 TensorFlow 中，我会喜欢：

with tf.device('/gpu:0'):
    An = matpow(A, n)
with tf.device('/gpu:1'):
    Bn = matpow(B, n)
with tf.Session() as sess:
    C = sess.run(An + Bn)

但是，由于 PyTorch 是动态的，所以我在做同样的事情时遇到了麻烦。我尝试了以下方法，但只需要更多时间。

with torch.cuda.device(0):
    A = A.cuda()       
with torch.cuda.device(1):
    B = B.cuda()
C = matpow(A, n) + matpow(B, n).cuda(0)

我知道有一个模块可以在批处理维度上并行化模型，torch.nn.DataParallel但在这里我尝试做一些更基本的事情。

4

1 回答 1

1

您可以为此使用 cuda流。这不一定将它分布在两个设备上，但执行将是并行的。

s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

with torch.cuda.stream(s1):
    A = torch.pow(A,n)
with torch.cuda.stream(s2):
    B = torch.pow(B,n)

C = A+B

虽然我不确定如果你只并行化这一个操作是否真的会加速你的计算。你的矩阵一定很大。

如果您的要求是跨设备拆分它，您可以在流之前添加它：

A = A.cuda(0)
B = B.cuda(1)

然后在电源操作后，您需要再次将它们放在同一设备上，例如B = B.cuda(0)。之后，您可以进行添加。

于 2018-01-31T11:14:11.137 回答