PyTorch中如何对4D张量做指定索引，该场景是否适用torch.gather？

是否适合使用torch.gather #

这个场景不适合使用torch.gather，和官方文档的规则要求一致：

• torch.gather仅支持沿单个指定维度收集输入张量的元素，要求索引张量和输入张量除了待收集的维度外，其余所有维度的形状必须完全匹配，不支持同时对多个维度做独立索引。
• 你的需求是同时对输入张量的4个维度分别指定索引取数，不属于单维度收集的适用场景，因此无需使用gather实现。

等价高效实现方案 #

直接使用PyTorch原生的高级整数索引即可，一行代码就能实现和你提供的循环逻辑完全一致的效果，且为向量化运算，执行效率远高于Python循环：

# 注意不要用input作为变量名，input是Python内置关键字，此处用input_tensor指代你的输入张量
output = input_tensor[idx[0], idx[1], idx[2], idx[3]]

执行后得到的output张量形状就是你期望的(1632,)。

补充说明：你提供的idx张量形状为(4, 1632)，拆分后idx[0]到idx[3]四个张量形状均为(1632,)，刚好对应1632个待取元素在四个维度上的索引值，符合PyTorch多维度整数索引的规则要求。

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