PyTorch矢量化求和不同于循环求和

dm7nw8vv  于 2022-11-29  发布在  其他
关注(0)|答案(1)|浏览(138)

我在使用torch 1.7.1时注意到,如果索引重复,向量化求和与循环中的求和是不同的。例如:

import torch

indices = torch.LongTensor([0,1,2,1])
values = torch.FloatTensor([1,1,2,2])
result = torch.FloatTensor([0,0,0])

looped_result = torch.zeros_like(result)

for i in range(indices.shape[0]):
    looped_result[indices[i]] += values[i]

result[indices] += values

print('result:',result)
print('looped result:', looped_result)

结果:

result tensor: ([1., 2., 2.])
 looped result tensor: ([1., 3., 2.])

正如你所看到的,循环变量有正确的和,而矢量化变量没有。有没有可能避免循环,仍然得到正确的结果?

f1tvaqid

f1tvaqid1#

这里的问题是你在同一个索引中多次索引result,这对于这个就地操作来说是注定要失败的。相反,你需要使用index_addindex_add_,例如(作为你的代码片段的延续):

>>> result_ia = torch.zeros_like(result)
>>> result_ia.index_add_(0, indices, values)
tensor([1., 3., 2.]

相关问题