Question

我有3个向量-和向量，贡献向量和值向量。我想根据值向量对它们的贡献向量求和，并将它们放在和向量中的相应索引中。一个例子是：

A = [0;0] (sum vector), B = [0,0,1,1] (contribution vector) C=[20,30,40,10] (value vector)

输出： A = [20+30;40+10]

B向量的长度与C相同，并且它们的相应索引告诉我们要在A中的哪个位置添加。

我能够通过如下的for循环来实现：

for index,value in enumerate(C):
    A[B[index]]+=value

但是，由于这将是我的NN模型正向循环的一部分，因此将导致严重的性能问题。具体来说，我在寻找一种矢量/矩阵排序方法，该方法将更加高效。在上面的示例中，对我而言有效的工作是：

A=torch.zeros(2,1)
C=C.reshape(2,2)
sum=torch.sum(C,1).reshape(2,1)
A += sum

但是，我遇到了问题，因为并非总是A的索引具有相同的贡献。例如-B = [0,0,0,1,1]和C=[20,30,40,10,50]的情况。对于一般情况，是否有功能或策略性的方法来做到这一点？谢谢！

Answer 1

执行循环会在某种程度上降低性能。假设贡献向量中唯一类型的数量远小于B（或C）的长度，则可以使用O(num_types) << O（len_B）`：

进行如下for循环


num_types = 3
B_len = 5
C_len = B_len

B = torch.randint(0, num_types, size=[B_len,])
"""
>>> B
tensor([2, 1, 1, 0, 0])
"""

C = torch.randint(0, 10, size=[C_len,])
C = C.float()
"""
>>> C
tensor([1., 5., 7., 6., 2.])
"""

# For loop here
A = [torch.sum(C * (torch.eq(B, type).float()) for type in range(num_types)]
# A = [tensor(8.), tensor(12.), tensor(1.)]

# Convert it to torch.tensor
A = torch.stack(A)
# tensor([ 8., 12.,  1.])

Answer 2

您正在寻找index_add_()

A.index_add_(0, B, C)

请注意，B的类型应为torch.long（它是索引向量），而C的类型应为torch.float，与{{ 1}}。
此外，在A和dim是多维张量的情况下，可以使用第一个A参数沿不同维度进行此求和。

如何基于索引向量求和

2 个答案: