PyTorch如何将值赋值给gather或topk获取的指定位置元素

最优实现方案 #

你要实现的是gather的反向赋值操作，PyTorch提供了原生的scatter_原位方法来完成这个需求，是当前性能最优、代码最简洁的实现：

import torch
# 你的原有代码
p = torch.randn(5, 7)
val, idx = p.topk(3, dim=-1)
x = torch.randn(5, 3)

# 核心赋值代码
p.scatter_(dim=-1, index=idx, src=x)

方法说明 #

scatter_是gather的逆操作，会直接修改原张量p的值：沿着dim=-1维度，把x中的每个元素，写入到p中idx指定的对应索引位置，完全匹配你要的赋值效果。

如果需要更显式的索引控制，也可以用高级索引方式实现，效果和性能和上述方案几乎无差异：

# 生成行索引，和列索引idx形状对齐
row_idx = torch.arange(p.shape[0]).unsqueeze(-1).repeat(1, idx.shape[1])
p[row_idx, idx] = x

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