张量索引与切片

访问和修改张量的特定部分是处理深度学习数据时常有的需求。无论您是需要选择单个数据点、提取一批训练样本、裁剪图像补丁，还是挑选特定特征，PyTorch都提供了强大且灵活的索引和切片机制，类似于NumPy数组中的那些，但PyTorch的机制与GPU加速和自动微分功能相集成。

基本索引

访问张量元素最直接的方式是使用标准的Python整数索引。请记住，PyTorch张量与Python列表和NumPy数组一样，使用0作为起始索引。

对于一维张量，您可以使用其索引访问元素：

import torch

# 创建一个一维张量
x_1d = torch.tensor([10, 11, 12, 13, 14])
print(f"原始一维张量:\n{x_1d}")

# 访问第一个元素
first_element = x_1d[0]
print(f"\n第一个元素 (x_1d[0]): {first_element}, 类型: {type(first_element)}")

# 访问最后一个元素
last_element = x_1d[-1]
print(f"最后一个元素 (x_1d[-1]): {last_element}")

# 修改一个元素
x_1d[1] = 110
print(f"\n修改后的张量:\n{x_1d}")

注意，访问单个元素会返回一个包含单个值的torch.Tensor（一个0维张量或标量），而不是一个标准的Python数字，除非您使用.item()明确提取它。元素的修改是原地进行的。

对于多维张量，您需要为每个维度提供索引，并用逗号分隔：

# 创建一个二维张量 (例如，一个小矩阵)
x_2d = torch.tensor([[1, 2, 3],
                     [4, 5, 6],
                     [7, 8, 9]])
print(f"原始二维张量:\n{x_2d}")

# 访问第0行第1列的元素
element_0_1 = x_2d[0, 1]
print(f"\n在 [0, 1] 的元素: {element_0_1}")

# 访问整个第一行 (索引0)
first_row = x_2d[0] # or x_2d[0, :]
print(f"\n第一行 (x_2d[0]): {first_row}")

# 访问整个第二列 (索引1)
second_col = x_2d[:, 1]
print(f"第二列 (x_2d[:, 1]): {second_col}")

# 修改一个元素
x_2d[1, 1] = 55
print(f"\n修改后的二维张量:\n{x_2d}")

提供的索引数量少于维数时，会沿剩余维度选择一个完整的子张量。例如，x_2d[0] 会选择整个第一行。

张量切片

切片允许您沿张量维度选择一系列元素。语法是 start:stop:step，其中 start 是包含的，stop 是不包含的，而 step 定义了间隔。省略 start 默认为0，省略 stop 默认为维度的末尾，省略 step 默认为1。

# 创建一个一维张量
y_1d = torch.arange(10) # Tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
print(f"原始一维张量: {y_1d}")

# 选择从索引2开始到（不包含）索引5的元素
slice1 = y_1d[2:5]
print(f"\n切片 y_1d[2:5]: {slice1}")

# 选择从开头到索引4的元素
slice2 = y_1d[:4]
print(f"切片 y_1d[:4]: {slice2}")

# 选择从索引6到末尾的元素
slice3 = y_1d[6:]
print(f"切片 y_1d[6:]: {slice3}")

# 选择每隔一个的元素
slice4 = y_1d[::2]
print(f"切片 y_1d[::2]: {slice4}")

# 选择从索引1到7的元素，步长为2
slice5 = y_1d[1:8:2]
print(f"切片 y_1d[1:8:2]: {slice5}")

# 反转张量
slice6 = y_1d[::-1]
print(f"切片 y_1d[::-1]: {slice6}")

切片对多维张量的工作方式类似。您可以将整数索引和切片结合使用：

# 创建一个3x4张量
x_2d = torch.tensor([[ 0,  1,  2,  3],
                     [ 4,  5,  6,  7],
                     [ 8,  9, 10, 11]])
print(f"原始二维张量:\n{x_2d}")

# 选择前两行以及第1和第2列
sub_tensor1 = x_2d[0:2, 1:3]
print(f"\n切片 x_2d[0:2, 1:3]:\n{sub_tensor1}")

# 选择所有行，但只选择最后两列
sub_tensor2 = x_2d[:, -2:]
print(f"\n切片 x_2d[:, -2:]:\n{sub_tensor2}")

# 选择第一行，从第1列到末尾
sub_tensor3 = x_2d[0, 1:]
print(f"\n切片 x_2d[0, 1:]:\n{sub_tensor3}")

# 选择第0行和第2行（使用步长），所有列
sub_tensor4 = x_2d[::2, :]
print(f"\n切片 x_2d[::2, :]:\n{sub_tensor4}")

张量 x_2d 使用 x_2d[0:2, 1:3] 进行切片的视觉表示。它选择第0和第1行，以及第1和第2列。

切片的一个重要特性（与某些其他形式的索引不同）是，返回的张量通常与原始张量共享底层存储。修改切片会修改原始张量。

print(f"修改切片前的原始 x_2d:\n{x_2d}")

# 获取一个切片
sub_tensor = x_2d[0:2, 1:3]

# 修改切片
sub_tensor[0, 0] = 101

print(f"\n修改后的切片:\n{sub_tensor}")
print(f"\n修改切片后的原始 x_2d:\n{x_2d}") # 注意变化！

如果您需要一个不共享内存的副本，请在切片上使用 .clone(): sub_tensor_copy = x_2d[0:2, 1:3].clone()。

布尔索引 (遮罩)

您可以使用布尔张量来索引另一个张量。布尔张量的形状必须能够广播到被索引张量的形状（通常，它们的形状完全相同）。只有布尔张量中对应 True 值的元素（即“遮罩”）才会被选中。这对于根据条件筛选数据非常有用。

# 创建一个张量
data = torch.tensor([[1, 2], [3, 4], [5, 6]])
print(f"原始数据张量:\n{data}")

# 创建一个布尔遮罩 (例如，选择大于3的元素)
mask = data > 3
print(f"\n布尔遮罩 (data > 3):\n{mask}")

# 应用遮罩
selected_elements = data[mask]
print(f"\n通过遮罩选择的元素:\n{selected_elements}")
print(f"所选元素的形状: {selected_elements.shape}")

# 根据条件修改元素
data[data <= 3] = 0
print(f"\n将小于等于3的元素设置为零后的数据:\n{data}")

布尔索引通常返回一个包含所有选定元素的一维张量。与切片不同，它不保留原始形状。此外，布尔索引通常会创建一个副本，而不是一个视图。

您可以将布尔索引与其他形式结合使用。例如，根据应用于其中一列的条件选择行：

# 选择第一列大于2的行
row_mask = data[:, 0] > 2
print(f"\n行遮罩 (data[:, 0] > 2): {row_mask}")

selected_rows = data[row_mask, :] # 使用 ':' 选择所选行中的所有列
# Or simply: data[row_mask] - PyTorch 通常会推断出完整的行选择
print(f"\n第一列大于2的行:\n{selected_rows}")

整数数组索引

除了单个整数和切片，您还可以使用列表或一维整数张量沿维度进行索引。这使得您可以按任意顺序选择元素，或多次选择相同的元素。

x = torch.arange(10, 20) # Tensor([10, 11, 12, 13, 14, 15, 16, 17, 18, 19])
print(f"原始一维张量: {x}")

indices = torch.tensor([0, 4, 2, 2]) # 注意索引2的重复
selected = x[indices]
print(f"\n使用索引 {indices} 选择的元素: {selected}")

# 对于二维张量
y = torch.arange(12).reshape(3, 4)
# [[ 0,  1,  2,  3],
#  [ 4,  5,  6,  7],
#  [ 8,  9, 10, 11]]
print(f"\n原始二维张量:\n{y}")

# 选择特定行
row_indices = torch.tensor([0, 2])
selected_rows = y[row_indices] # 选择第0行和第2行
print(f"\n使用索引 {row_indices} 选择的行:\n{selected_rows}")

# 选择特定列
col_indices = torch.tensor([1, 3])
selected_cols = y[:, col_indices] # 从所有行中选择第1列和第3列
print(f"\n使用索引 {col_indices} 选择的列:\n{selected_cols}")

# 使用索引对选择特定元素
row_idx = torch.tensor([0, 1, 2])
col_idx = torch.tensor([1, 3, 0])
selected_elements = y[row_idx, col_idx] # 选择 (0,1), (1,3), (2,0) -> [1, 7, 8]
print(f"\n使用 (row_idx, col_idx) 选择的特定元素:\n{selected_elements}")

与布尔索引类似，整数数组索引通常返回一个新的张量（一个副本），而不是原始张量存储的视图。输出的形状取决于索引方法。当选择完整的行或列时，其他维度会被保留。当为多个维度提供索引数组（例如 y[row_idx, col_idx]）时，结果通常是一个对应于所选元素的一维张量。

掌握这些索引和切片技术能够精准地控制张量数据，为后续步骤中的数据准备、特征提取以及模型输入输出的操作奠定根基。