如何在PyTorch中用2D索引张量从3D张量提取行？

要实现你要的效果，关键是让索引张量的形状匹配gather()的要求——除了指定的操作维度外，其余维度需要和输入张量兼容（可广播）。具体步骤如下：

1. 调整索引张量形状：
   你的输入张量x形状是(2,3,3)（i=2，j=3，k=3），索引index是(1,3)。我们需要：

   • 把index在i维度（第0维）扩展为2个样本的索引，让每个i样本都用同一组j维度索引；
   • 在最后添加一个维度（适配k维度），这样gather()会自动提取对应j行的所有k元素（完整行）。

2. 执行gather操作：
   指定dim=1（对应j维度），用调整后的索引张量提取数据。

完整代码示例：

import torch

x = torch.tensor([[[1,2,3],
                   [4,5,6],
                   [7,8,9]],

                  [[10,11,12],
                   [13,14,15],
                   [16,17,18]]])
index = torch.tensor([[1,2,0]])

# 调整索引：扩展到i维度的2个样本，再添加k维度的适配维度
adjusted_index = index.expand(2, 3).unsqueeze(-1)
# 执行gather，dim=1表示在j维度按索引取值
output = x.gather(dim=1, index=adjusted_index)

print(output)

输出结果：

tensor([[[4, 5, 6],
         [7, 8, 9],
         [1, 2, 3]],

        [[13, 14, 15],
         [16, 17, 18],
         [10, 11, 12]]])

关键说明： #

• index.expand(2,3)：不复制数据，高效地将原索引扩展为(2,3)，让两个i样本共用同一组j索引；
• unsqueeze(-1)：将索引形状变为(2,3,1)，这样gather()在dim=1操作时，会把每个j索引应用到k维度的所有元素，从而提取完整的行；
• 如果用index.repeat(2,1)代替expand也能实现，但expand更节省内存。

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