torch.gather(...) 调用的设置结果答案

【问题标题】：Setting results of torch.gather(...) callstorch.gather(...) 调用的设置结果
【发布时间】：2021-11-05 01:32:28
【问题描述】：

我有一个形状为 n x m 的 2D pytorch 张量。我想使用索引列表（可以使用 torch.gather 完成）对第二个维度进行索引，然后然后还设置新值到索引的结果。

例子：

data = torch.tensor([[0,1,2], [3,4,5], [6,7,8]]) # shape (3,3)
indices = torch.tensor([1,2,1], dtype=torch.long).unsqueeze(-1) # shape (3,1)
# data tensor:
# tensor([[0, 1, 2],
#         [3, 4, 5],
#         [6, 7, 8]])

我想为每行选择指定的索引（这将是 [1,5,7]，但随后也将这些值设置为另一个数字 - 例如 42

我可以通过以下方式逐行选择所需的列：

data.gather(1, indices)
tensor([[1],
        [5],
        [7]])
data.gather(1, indices)[:] = 42 # **This does NOT work**, since the result of gather 
                                # does not use the same storage as the original tensor

这很好，但我现在想更改这些值，并且更改也会影响data 张量。

我可以用它做我想做的事，但它似乎很不符合pythonic：

max_index = torch.max(indices)
for i in range(0, max_index + 1):
  mask = (indices == i).nonzero(as_tuple=True)[0]
  data[mask, i] = 42
print(data)
# tensor([[ 0, 42,  2],
#         [ 3,  4, 42],
#         [ 6, 42,  8]])

关于如何更优雅地做到这一点的任何提示？

【问题讨论】：

标签： python indexing pytorch tensor

【解决方案1】：

您要查找的是带有value 选项的torch.scatter_。

Tensor.scatter_(dim, index, src, reduce=None) → Tensor
将张量 src 中的所有值写入 self 在 index 张量中指定的索引处。对于src 中的每个值，其输出index 指定为它在 src 中的索引为dimension != dim 和相应的值在dimension = dim 的索引中。

以二维张量为输入，dim=1，运算为：
self[i][index[i][j]] = src[i][j]

_{虽然没有提到 value 参数...}

使用value=42 和dim=1，这将对数据产生以下影响：

data[i][index[i][j]] = 42

这里就地应用：

>>> data.scatter_(index=indices, dim=1, value=42)
>>> data
tensor([[ 0, 42,  2],
        [ 3,  4, 42],
        [ 6, 42,  8]])

【讨论】：