如何基于索引向量求和

小偷

我有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]对于一般情况,是否有功能或策略性的方法来做到这一点?谢谢!

hai

你在找 index_add_()

A.index_add_(0, B, C)

请注意,该B类型应为torch.long(它是一个索引向量),并且C应类型torch.float与相同A
此外,还可以使用第一个dim参数来做到沿着不同的维度这个总和的情况下,AC是多维张量。

本文收集自互联网,转载请注明来源。

如有侵权,请联系 [email protected] 删除。

编辑于
0

我来说两句

0 条评论
登录 后参与评论

相关文章