进阶必看，3种灵活操作PyTorch张量的高级方法

大家好，在PyTorch中进行高级张量操作时，开发者经常面临这样的问题，如何根据一个索引张量从另一个张量中选取元素。

例如有一个包含数千个特征的大规模数据集，需要根据特定的索引模式快速提取信息。本文将介绍三种索引选择方法来解决这类问题。

1 torch.index_select

torch.index_select函数通过在指定的维度上进行元素选择，同时在其他维度上保持元素不变。也就是说，在目标维度上根据索引张量来挑选元素，而其他维度的元素则原封不动。为了更直观地理解这一概念，来看一个2D张量的示例，这里将沿着维度1进行元素的选择：

num_picks = 2

values = torch.rand((len_dim_0, len_dim_1))
indices = torch.randint(0, len_dim_1, size=(num_picks,))
# [len_dim_0, num_picks]
picked = torch.index_select(values, 1, indices)

由此得到的张量形状为[len_dim_0, num_picks]：对于维度0上的每个元素，都从维度1中选取了相同的元素。将其形象化：

现在迈入三维张量的世界，这样更贴近机器学习与数据科学的实际需求。

设想一个三维张量，其维度为[batch_size, num_elements, num_features]