我有一个大向量需要更新。我将通过向向量中的特定元素添加偏移量来更新它。我指定要更新的索引向量(调用索引向量 ix
),并为每个索引指定一个要添加到该元素的值(调用值向量 >vals
).如果索引向量的所有条目都是唯一的,则以下代码就足够了:
vec = torch.zeros(4, dtype=torch.float)
ix = torch.tensor([0,2], dtype=torch.long)
vals = torch.tensor([0.2, 0.5], dtype=torch.float)
vec[ix] += vals
但是,如果 ix
中有重复的索引,这将不起作用。对于重复索引的情况,一种天真的方法如下:
for i in range(len(ix)):
vec[ix[i]] += vals[i]
但这不能很好地扩展 - 当 ix
很大时它会很慢。有没有更快的方法来做到这一点?如果有一种快速方法可以对 ix
中具有相同索引的所有 vals
条目求和,那么解决方案应该很简单。
更新:
我找到了一种效果很好的解决方案,如下所述。我仍然希望获得更好的解决方案的反馈。
# get unique indices
ix_unique = torch.unique(ix)
# for each unique index, get sum of all vals with that index
vals_unique = torch.stack([
torch.sum(torch.where(ix==i, vals, torch.zeros_like(vals)))
for i in ix_unique
])
# update vec
vec[ix_unique] += vals_unique
最佳答案
如果您希望允许对同一个 ix 索引进行多次更新,还存在一个名为 pytorch_scatter 的库. 在这种情况下,例如ix 中的 3 个相同的索引将导致 3*val 被添加到该索引。
https://stackoverflow.com/questions/53321844/