我有一个张量列表，所有张量都存在于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]。因此，我想从张量res获得张量A，以使res的第一行是A的前5行的和（因为splitter[0] = 5）沿dim=0。 res的第二行是A的后3行（第5行到第7行）的总和。等等。我可以不使用for循环来做到这一点吗？或者我可以将此for循环并行化，因为它的操作彼此独立并且互斥且详尽无遗。
我希望添加的详细信息足够。如果我需要在问题中添加更多详细信息，请告诉我。在此先感谢：）
答案 0 :(得分：0)
如果拆分可以相同，则可以通过矢量化方法解决：
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]])
这里需要进行内存/速度折衷。希望这更接近您想要的。
答案 1 :(得分：0)
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