矩阵乘法的基本多GPU并行化

Question

我想并行上2个GPU简单下面的表达式:C = A^n + B^n通过计算A^n在GPU 0和B^n对结果求和之前对GPU 1.

在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但在这里我尝试做一些更基本的事情.

Answer 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)。之后就可以进行添加了。