如果循环中涉及的所有张量都在 GPU 上，我的 for 循环是不是并行运行？

Posted 2023-03-12

【中文标题】如果循环中涉及的所有张量都在 GPU 上，我的 for 循环是不是并行运行？

【英文标题】：Does my for loop run parallely, if all the tensors involved in the loop are on the GPU?如果循环中涉及的所有张量都在 GPU 上，我的 for 循环是否并行运行？

【发布时间】：2020-11-08 00:25:39

【问题描述】：

我有一个张量列表，它们都存在于 GPU 上。我通过使用torch.split 在 GPU 上拆分一个张量来获得此列表。我想获取我拥有的张量列表的总和列表。所以，简单来说，我想得到一个列表，其中第一个元素是列表中第一个张量的总和，依此类推。如果我为此运行一个 for 循环，它会并行化吗？如果没有，有没有办法让它并行运行？我想并行化它，因为列表很长，并且求和运算可以并行完成，并且独立于列表中存在的每个张量。如果可以在 GPU 上执行此操作，性能提升将是巨大的。

更新：考虑我有一个张量列表如下：

ls 
[tensor([[0.8469, 0.3712, 0.2956],
         [0.6548, 0.5284, 0.8682],
         [0.5748, 0.2390, 0.1402],
         [0.0010, 0.1794, 0.6048],
         [0.4636, 0.4101, 0.6543]], device='cuda:0'),
 tensor([[0.2138, 0.3613, 0.8712],
         [0.4689, 0.0503, 0.7342],
         [0.1368, 0.0688, 0.9223]], device='cuda:0'),
 tensor([[0.3131, 0.6142, 0.1555],
         [0.4099, 0.5000, 0.7578],
         [0.7353, 0.2425, 0.4407],
         [0.5943, 0.0377, 0.4820],
         [0.5898, 0.9585, 0.6993]], device='cuda:0'),
 tensor([[0.8629, 0.3172, 0.4248],
         [0.9957, 0.6998, 0.0931],
         [0.0258, 0.9898, 0.5250]], device='cuda:0'),
 tensor([[0.0298, 0.4033, 0.9465],
         [0.2763, 0.9412, 0.4873]], device='cuda:0')]

如您所见，我有一个包含 5 个不同形状的张量的列表。每个张量的第一维形状为 3。由于第 0 维，形状不同。所以，在这个例子中，列表中张量的形状是[[5,3], [3, 3], [5, 3], [3, 3], [2,3]]。我想从此列表中获取张量列表，如下所示：

sums = [torch.sum(li, axis=0) for li in ls]
sums
[tensor([2.5412, 1.7280, 2.5632], device='cuda:0'),
 tensor([0.8195, 0.4804, 2.5277], device='cuda:0'),
 tensor([2.6424, 2.3528, 2.5352], device='cuda:0'),
 tensor([1.8844, 2.0068, 1.0429], device='cuda:0'),
 tensor([0.3062, 1.3445, 1.4338], device='cuda:0')]

因此，如您所见，列表中的第一个张量是列表中的第一个张量 ls 沿维度 0 的总和。第二个张量是列表ls 中第二个张量沿维度0 之和，依此类推。

为了完成这项任务，我目前正在使用 for 循环。它迭代地计算总和并将其附加到sums 列表中。但是，这是非常低效的，因为我的张量列表非常大，大约 100K，并且在每次迭代中这样做非常低效。我想知道是否有任何方法可以更有效地做到这一点。

张量列表ls 是通过像这样拆分一个大张量获得的：

splitter = [5, 3, 5, 3, 2]

A = torch.rand(18, 3).cuda()

ls = torch.split(A, splitter)
ls
(tensor([[0.1969, 0.6113, 0.3563],
         [0.9180, 0.7759, 0.5953],
         [0.0279, 0.4014, 0.2268],
         [0.9026, 0.3821, 0.1498],
         [0.3630, 0.9144, 0.3277]], device='cuda:0'),
 tensor([[2.1312e-02, 5.2311e-01, 8.9177e-02],
         [4.7427e-01, 2.4503e-04, 1.2559e-01],
         [5.1641e-01, 9.1357e-01, 9.5637e-01]], device='cuda:0'),
 tensor([[0.3730, 0.4251, 0.9437],
         [0.5634, 0.3086, 0.5891],
         [0.5602, 0.0872, 0.2128],
         [0.7717, 0.1920, 0.3977],
         [0.5787, 0.3488, 0.7499]], device='cuda:0'),
 tensor([[0.9338, 0.4330, 0.8843],
         [0.5646, 0.0574, 0.8790],
         [0.4692, 0.5831, 0.9160]], device='cuda:0'),
 tensor([[0.9786, 0.5209, 0.9364],
         [0.4370, 0.4917, 0.3672]], device='cuda:0'))

因此，如果无法避免 for 循环，那么根据提供的拆分器，是否有人对对主张量 A 求和有任何想法？因此，例如，在上面的代码中，拆分器是[5, 3, 5, 3, 2]。所以，我想从张量A 中获得一个张量res，这样res 的第一行是A（因为splitter[0] = 5）沿dim=0 的前5 行的总和。 res 的第二行是 A 的下 3 行（第 5 行到第 7 行）的总和。等等。我可以在不使用 for 循环的情况下执行此操作吗？或者我可以并行化这个 for 循环，因为它正在执行的操作是相互独立的，并且是相互排斥和详尽的。

我希望添加的细节就足够了。如果我需要在问题中添加任何进一步的细节，请告诉我。在此先感谢:)

【问题讨论】：

您能否在此处添加一些代码，以便我们更好地了解您在做什么以及您想要实现什么？

@VictorZuanazzi，你能检查更新的问题吗？谢谢。

【参考方案1】：

如果拆分可以相同，那么你可以用矢量化的方式来解决它：

splitter = [6, 6, 6]

A = torch.rand(18, 3).cuda()

A_splits = A.reshape(-1, len(splitter), 3)

sums = A_splits.sum(dim=1)

这不是您正在寻找的通用解决方案，但也许它已经解决了您的问题？

编辑：

理想情况下，您可以用矢量化操作（例如.sum(dim=1)）替换循环，但矢量化操作仅适用于张量数据。如果张量之间的差异不是很大，您可以使用零将所有张量填充为相同的形状。

splitter = [5, 3, 5, 3, 2] # largest number of tensors is 5

A = torch.rand(18, 3).cuda()

A_pad = torch.zeros(max(splitter) * len(splitter), 3)

splitter_index = torch.tensor([i +  (max(splitter) * n) for n, l in enumerate(splitter) for i in range(l)])

A_pad[splitter_index] =  A

A_sum = A_pad.view(-1, max(splitter), 3).sum(dim=1) # double check the dim

A_sum

tensor([[2.2903, 2.3379, 2.6550],
        [1.1394, 1.2519, 0.7374],
        [1.7970, 2.8287, 2.4855],
        [0.7964, 1.1991, 1.4032],
        [1.8656, 0.4916, 0.2935]])

这里需要权衡内存/速度。希望这更接近您正在寻找的内容。

【讨论】：

不是真的，我的分裂总是不同的。实际上，拆分器只是因为可变大小的训练数据的串联。最后，我使用这个分割器来分割向量化的训练数据。

【参考方案2】：

PyTorch 异步运行 GPU 操作 (see docs)。

当您调用使用 GPU 的函数时，操作会被排入特定设备的队列

这意味着，您的求和运算可以并行运行。

我做了一个简单的实验来测试这个。如果我是对的，那就证明你不用担心这里的并行性。

import torch

A = torch.rand(100000, 32, device='cuda')
splits = torch.split(A, 4)

您的代码：

%%timeit -r1 -n5
sums = [s.sum() for s in splits]
torch.cuda.synchronize()

# Output: 5 loops, best of 1: 374 ms per loop

在每次求和运算后添加同步：

%%timeit -r1 -n5
sums = [torch.cuda.synchronize() or s.sum() for s in splits]

# Output: 5 loops, best of 1: 897 ms per loop