如何根据索引向量进行求和

我有三个向量 – 一个求和向量，一个贡献向量和一个值向量。我希望根据贡献向量对值向量进行求和，并将结果放置在求和向量的相应索引中。一个例子是：

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

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

其中B向量与C向量的长度相同，它们的对应索引告诉我们将C的值加到A的哪个位置上。

我通过使用for循环实现了这个功能，如下所示：

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

然而，由于这将是我神经网络模型的前向循环的一部分，它会引起显著的性能问题。特别是我在寻找一种更高效的向量/矩阵排序方法。在上面的例子中，对我有效的方法是：

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]时。对于一般情况，是否有函数或策略可以实现这个功能？谢谢！

回答：

你正在寻找index_add_()

A.index_add_(0, B, C)

请注意，B应当是torch.long类型（它是一个索引向量），而C应当是torch.float类型，与A的类型相同。
此外，如果A和C是多维张量，你可以使用第一个dim参数来沿不同维度进行这种求和操作。

学技术