如何根据给定索引矩阵为Tensor矩阵赋值（无循环实现）

给定以下PyTorch Tensor：

import torch

index = torch.tensor([[0, 1],
                      [1, 2],
                      [4, 1]])
a = torch.zeros((3, 6), dtype=torch.float32)
b = torch.tensor([5., 6., 7., 8., 9., 10.])

需要将b中对应index的值赋值给a的指定位置，无需循环即可实现，以下是两种正确的实现方式：

方法一：使用torch.Tensor.scatter_ #

scatter_适合在指定维度上，将源张量的值根据索引分散到目标张量中。这里我们在列维度（dim=1）操作：

# 生成与index形状匹配的源张量：从b中取出index对应的元素
src = b[index]
# 原地修改a，将src的值放入index指定的列位置
a.scatter_(dim=1, index=index, src=src)

执行后a的结果符合预期：

tensor([[5., 6., 0., 0., 0., 0.],
        [0., 6., 7., 0., 0., 0.],
        [0., 6., 0., 0., 9., 0.]])

方法二：使用torch.Tensor.index_put_ #

index_put_需要传入行、列两个维度的索引，精准定位赋值位置：

# 生成行索引：每一行对应index的两个元素
row_indices = torch.arange(index.size(0)).unsqueeze(1).expand_as(index)
# 取出要赋值的目标值
values = b[index]
# 原地修改a，将values放入(row_indices, index)指定的位置
a.index_put_((row_indices, index), values)

同样可以得到预期结果。

关键注意点 #

• 两种方法都是无循环的向量操作，符合PyTorch高效计算的特性。
• 必须保证src/values的形状与index完全匹配，才能正确对应到目标位置。

内容的提问来源于stack exchange，提问作者Young By