我在使用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.])
正如你所看到的,循环变量有正确的和,而矢量化变量没有。有没有可能避免循环,仍然得到正确的结果?
1条答案
按热度按时间f1tvaqid1#
这里的问题是你在同一个索引中多次索引
result
,这对于这个就地操作来说是注定要失败的。相反,你需要使用index_add
或index_add_
,例如(作为你的代码片段的延续):