如何通过索引张量从3D PyTorch张量中选取特定行？

针对你给出的场景，这里提供两种简洁的实现方式：

方法一：使用高级索引 #

利用PyTorch的高级索引特性，直接按维度匹配选取元素：

import torch

a = torch.FloatTensor([[[1,1,1],[2,2,2]],[[9,9,9],[5,5,5]]])
b = torch.IntTensor([1,0])

# 第一维遍历所有batch，第二维选取b中对应索引的行
result = a[torch.arange(a.size(0)), b]
print(result)

输出结果：

tensor([[2., 2., 2.],
        [9., 9., 9.]])

方法二：使用torch.gather()函数 #

gather函数可以在指定维度上按索引提取元素，需要先调整索引张量的形状以匹配源张量的维度结构：

import torch

a = torch.FloatTensor([[[1,1,1],[2,2,2]],[[9,9,9],[5,5,5]]])
b = torch.IntTensor([1,0])

# 将索引张量扩展为(2,1,1)，与源张量除dim=1外的维度一致
b_expanded = b.unsqueeze(1).unsqueeze(2)
# 在dim=1维度上按索引提取，再压缩多余维度
result = a.gather(dim=1, index=b_expanded).squeeze(1)
print(result)

同样会得到符合期望的输出。

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