张量的索引

 

整数列表选择

x = torch.tensor([[[10, 20, 30], [40, 50, 60]]])
x1 = x[[0, 0]]
x2 = x[[0, 0], [0, 1]]
x3 = x[0, [0, 1], [0, 1]]
print(x1)
print(x2)
print(x3)

将输出：

tensor([[[10, 20, 30],
[40, 50, 60]],

[[10, 20, 30],
[40, 50, 60]]])
tensor([[10, 20, 30],
[40, 50, 60]])
tensor([10, 50])

 

x1 = x[[0, 0]] 作用在第0维度，列表[0,0]，这意味着在维0上选择第0个元素两次
x2 分别作用在维0和维1上，选择两个元素， 分别是0维上的 0 和 1维度上的 0  与0维度上的0和1维度上的1 
x3 作用在三个维度，分别是0维上的 0 和 1维度上的 0 和2维度上的0 与 0维度上的0和1维度上的1和2维度上的1

布尔列表选择

x4 = torch.tensor([[10, 20, 30], [40, 50, 60]])
x5 = x4[[False, True], [False, True, True]]

tensor([50, 60])

x5作用在两个维度上，布尔列表作为索引， 值选择对应true的位置

 

张量作为索引

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
x2 = x[torch.tensor([[0, 1], [1, 0]])]
print(f"x2:{x2}")

x2:tensor([[[10, 20, 30],
[40, 50, 60]],

[[40, 50, 60],
[10, 20, 30]]])

选择x的第一个维度，

 

x = torch.tensor([[10, 20, 30], [40, 50, 60]])

col_idx = torch.tensor([0, 2])  # 选择第 0 列和第 2 列

result = x[:, col_idx]  # 选择所有行的第 0 列和第 2 列
print(result)
# 输出: tensor([[10, 30],
#              [40, 60]])

 

布尔张量作为索引

import torch

x = torch.tensor([[10, 20, 30], [40, 50, 60]])  # 形状 (2, 3)
index = torch.tensor([[False, True, True], [False, True, True]])  # 形状 (2, 3)
result = x[index]

tensor([20, 30, 50, 60])

包含 x 中所有 index 对应位置为 True 的元素，按 ​行优先顺序 排列

 

x = torch.tensor([[10, 20, 30], [40, 50, 60]])
index = torch.tensor([False, True])
result = x[index]
print(result)

tensor([[40, 50, 60]])

布尔张量作用于张量的第一个维度

布尔张量作为索引的关键在于先确定布尔张量作用在哪个维度