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

【问题标题】：Basic multi GPU parallelization of matrix multiplication矩阵乘法的基本多 GPU 并行化
【发布时间】：2017-11-06 09:19:46
【问题描述】：

我想在 2 个 GPU 上并行化以下简单表达式：C = A^n + B^n，方法是在对结果求和之前计算 GPU 0 上的 A^n 和 GPU 1 上的 B^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 在批处理维度上并行化模型，但在这里我尝试做一些更基本的事情。

【问题讨论】：

标签： python gpu pytorch

【解决方案1】：

您可以为此使用 cuda streams。这不一定会在两个设备上分发它，但会并行执行。

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

【讨论】：