如何根据指标求和

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模型前向循环的一部分，这将导致严重的性能问题。特别是我在寻找一种更有效的向量/矩阵排序方法。在上面的例子中，对我有效的方法是：

^{pr2}$

但是，我遇到了一些问题，因为A的索引并不总是具有相同的贡献率。例如-B = [0,0,0,1,1]和{}的情况。在一般情况下，是否有一种功能或战略方法来做到这一点？谢谢！在

Answer 1

执行for循环会在一定程度上降低性能。假设贡献向量中唯一类型的数量远小于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

您正在寻找^{}

A.index_add_(0, B, C)

注意，B应该是torch.long（它是一个索引向量），而{}应该是{}类型，与{}的类型相同。
此外，如果A和{}是多维张量，则可以使用第一个dim参数沿不同维度进行求和。在