pytorch 使用索引列表修改数组

Question

假设我有一个索引列表，并希望使用该列表修改现有数组。目前我能做到这一点的唯一方法是使用如下的 for 循环。只是想知道是否有更快/更有效的方法。

torch.manual_seed(0)
a = torch.randn(5,3)
idx = torch.Tensor([[1,2], [3,2]], dtype=torch.long)
for i,j in idx:
    a[i,j] = 1

我最初认为gather 或index_select 会在某种程度上回答这个问题，但看看documentation 这似乎不是答案。

在我的特定情况下，a 是一个 5 维向量，idx 是一个 Nx5 向量。所以我期望的输出（在使用a[idx] 之类的下标之后）是一个(N,) 形状的向量。

回答

感谢下面的@shai，我正在寻找的答案是： a[idx.t().chunk(chunks=2,dim=0)]。取自SO answer。

Answer 1

很简单

a[idx[:,0], idx[:,1]] = 1

您可以在this thread找到更通用的解决方案。