如何矢量化总和聚合？答案

【问题标题】：How to vectorize the sum aggregation?如何矢量化总和聚合？
【发布时间】：2021-12-05 23:25:22
【问题描述】：

我试图对这段 Python（或经过一些小改动后的 Matlab）代码进行矢量化，这是有向图的总和聚合

for j in range(batchSize):
    for i in range(2*nEdges[j]):
        localSum[j,receivers[j,i],0:2]+=(localFeature[j,i,0:2])

当它是= 时，我知道如何删除其中一个循环，但是我没有找到一种方法来处理这个+= 案例。做同样的技巧的困难在于receivers[j,i]对于不同的j和is可能是相同的值。

有人知道在这种情况下如何进行矢量化吗？

非常感谢:)

【问题讨论】：

【解决方案1】：

终于找到了。 torch_scatter 包提供了这个功能（对于 pytorch 张量）。嵌套循环可以重写为

localSum = torch_scatter.scatter(localFeature,receivers,dim=1,reduce='sum')

【讨论】：