python - 如何根据索引向量求和

Question

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

A = [0;0](和向量)，B = [0,0,1,1](贡献向量)C=[20,30,40,10](值向量)

输出: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

您正在寻找 index_add_()

A.index_add_(0, B, C)

请注意B类型应为 torch.long (它是一个索引向量)，和 C类型应为 torch.float ，与 A 的类型相同.
而且，您可以使用第一个dim在案例 A 中沿不同维度进行求和的参数和C是多维张量。