如何在PyTorch中用新矩阵替换原矩阵的指定元素？

以下两种方法均无需循环，符合你的需求：

方法一：利用torch.put_实现 #

put_可通过线性索引直接将源张量的值赋值到目标张量的指定位置，步骤如下：

import torch

a = torch.Tensor([[1,2,3],[4,5,6], [7,8,9]])
b = torch.tensor([[1,1],[1,1]], dtype=torch.float32)  # 确保与a的数据类型一致

# 定义目标位置的行、列索引
target_rows = [0, 0, 2, 2]
target_cols = [0, 2, 0, 2]
# 转换为行优先的线性索引
linear_indices = torch.tensor(target_rows) * a.size(1) + torch.tensor(target_cols)

# 将b展平后赋值到a的对应位置
a.put_(linear_indices, b.flatten())

# 查看结果
print(a)

执行后输出：

tensor([[1., 2., 1.],
        [4., 5., 6.],
        [1., 8., 1.]])

方法二：利用torch.scatter_实现 #

通过构造匹配形状的索引矩阵，使用scatter_完成分散赋值：

import torch

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

# 构造与b形状一致的行、列索引矩阵
row_indices = torch.tensor([[0, 0], [2, 2]])
col_indices = torch.tensor([[0, 2], [0, 2]])

# 提取目标行后按列索引进行scatter赋值
a[row_indices.flatten()].scatter_(dim=1, index=col_indices, src=b)

print(a)

同样可以得到目标结果。

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