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

Question

我有一个张量列表，所有这些都存在于 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张量是通过像这样 split 一个大张量来获得的:

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] .所以，我想获得一个张量 res来自张量 A这样第一行 res是 A 的前 5 行的总和(因为 splitter[0] = 5)沿 dim=0 .第二排res是 A 的下 3 行(第 5 行到第 7 行)的总和.等等。我可以在不使用 for 循环的情况下执行此操作吗？或者我可以并行化这个 for 循环，因为它正在做的操作是相互独立的，并且是相互排斥和详尽的。
我希望添加的细节足够了。如果我需要在问题中添加任何进一步的细节，请告诉我。提前致谢 :)

Answer 1

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

When you call a function that uses the GPU, the operations are enqueued to the particular device

这意味着，您的求和操作可能会并行运行。
我做了一个简单的实验来测试这一点。如果我是对的，那就证明你不需要担心这里的并行性。

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