【PyTorch】高级索引实战：从基础到多维组合索引的深度解析

1. 为什么你需要掌握PyTorch高级索引？

在深度学习项目中，我们经常需要处理各种形状的张量数据。想象一下，你正在处理一个图像分类任务，数据集包含50000张224x224像素的RGB图片，存储为一个形状为(50000, 3, 224, 224)的四维张量。这时候，如果你只想选择其中第100到200张图片的红色通道数据，或者筛选出所有包含猫的图片，该怎么办？

这就是PyTorch高级索引大显身手的时候了。与基本索引（如简单的切片操作）相比，高级索引提供了更灵活的数据访问方式。它允许你：

• 按任意顺序选择元素
• 重复选择相同元素
• 根据复杂条件筛选数据
• 高效处理高维数据

我在处理一个自然语言处理项目时就深有体会。当时需要从词嵌入矩阵中批量提取特定位置的词向量，如果没有掌握高级索引技巧，代码会变得冗长且低效。而使用高级索引后，原本需要十几行循环的代码，一行就能搞定。

2. 整数数组索引：精准定位数据元素

2.1 基础用法

整数数组索引是最直接的高级索引方式。它允许我们使用整数列表或张量来指定要访问的元素位置。让我们从一个简单的二维张量开始：

import torch

# 创建一个3x3的矩阵
x = torch.tensor([[10, 11, 12],
                  [13, 14, 15], 
                  [16, 17, 18]])
print("原始张量:\n", x)

# 选择第0行和第2行
rows = torch.tensor([0, 2])
print("\n选择的行:\n", x[rows])

# 选择第1列和第2列
cols = torch.tensor([1, 2])
print("\n选择的列:\n", x[:, cols])

输出结果：

原始张量:
 tensor([[10, 11, 12],
        [13, 14, 15],
        [16, 17, 18]])

选择的行:
 tensor([[10, 11, 12],
        [16, 17, 18]])

选择的列:
 tensor([[11, 12],
        [14, 15],
        [17, 18]])

2.2 高级技巧：组合索引

整数数组索引的真正威力在于多维组合使用。我们可以同时指定行和列的索引来精确定位元素：

# 选择(0,1)、(1,2)、(2,0)三个位置的元素
row_indices = torch.tensor([0, 1, 2])
col_indices = torch.tensor([1, 2, 0])
print("\n组合索引结果:\n", x[row_indices, col_indices])

输出：

组合索引结果:
 tensor([11, 15, 16])

这里有个重要细节：当使用多个索引数组时，PyTorch会将这些数组视为一组坐标对。也就是说，它会取(row_indices[0], col_indices[0])、(row_indices[1], col_indices[1])等位置的元素。

2.3 实际应用案例

在图像处理中，我们经常需要提取特定位置的像素值。假设我们有一个批量处理的图像张