python - 更新 PyTorch 中的特定向量元素

Question

我有一个大向量需要更新。我将通过向向量中的特定元素添加偏移量来更新它。我指定要更新的索引向量(调用索引向量 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

Answer 1

如果您希望允许对同一个 ix 索引进行多次更新，还存在一个名为 pytorch_scatter 的库.在这种情况下，例如ix 中的 3 个相同的索引将导致 3*val 被添加到该索引。