一般索引
从前往后提供的索引,会依次在每个维度上做索引。
import torch
a = torch.rand(4, 3, 28, 28)
print(a[0].shape)
print(a[0, 0].shape)
print(a[0, 0, 2, 4]) # 具体到某个元素
运行结果:
torch.Size([3, 28, 28])
torch.Size([28, 28])
tensor(0.3690)
普通的切片索引
print(a[:, :, 0:28:2, 0:28:2].shape) # step=2隔行采样
print(a[:, :, ::2, ::2].shape) # 等同于这个
运行结果:
torch.Size([4, 3, 14, 14])
torch.Size([4, 3, 14, 14])
注意负值的索引即表示倒数第几个元素,-2就是倒数第二个元素。
使用step
选择特定下标
选择特定下标有时候很有用,比如上面的a这个Tensor可以看作4张RGB(3通道)的MNIST图像,长宽都是28px。那么在第一维度上可以选择特定的图片,在第二维度上选择特定的通道。
# 选择第一张和第三张图
print(a.index_select(0, torch.tensor([0, 2])).shape)
# 选择R通道和B通道
print(a.index_select(1, torch.tensor([0, 2])).shape)
运行结果:
torch.Size([2, 3, 28, 28])
torch.Size([4, 2, 28, 28])
使用 …
在索引中使用表示任意多的维度
print(a[:, 1, ...].shape)
print(a[..., :2].shape)
print(a[0, ..., ::2].shape) # ::2表示间隔为2
运行结果:
torch.Size([4, 28, 28])
torch.Size([4, 3, 28, 2])
torch.Size([3, 28, 14]) # ::2表示间隔为2,原来最后一个维度值为28,因为取了间隔为2,所以现在最后一维的维度值变为了14
举例:
import numpy as np
a = np.arange(0, 6)
a = np.reshape(a, [2, 3])
print(a)
print(a[..., 1])
print(a[..., 0:1])
# 输出
[[0 1 2]
[3 4 5]]
[1 4]
[[0]
[3]]
## a的shape是(2, 3),a[…,1]就是输出最后一个维度的序号1的数值,
## 即是[1, 4],a[…, 0:1]保留最后一个维度0:1的数值。
b = np.arange(0, 24)
b = np.reshape(b, [2, 3, 4])
print(b)
print(b[..., 1])
print(b[..., 1:2])
# 输出
[[[ 0 1 2 3]
[ 4 5 6 7]
[ 8 9 10 11]]
[[12 13 14 15]
[16 17 18 19]
[20 21 22 23]]]
[[ 1 5 9]
[13 17 21]]
[[[ 1]
[ 5]
[ 9]]
[[13]
[17]
[21]]]
mask索引
可以获取满足一些条件的值的位置索引,然后用这个索引去取出这些位置的元素。
import torch
# 取出a这个Tensor中大于0.5的元素
a = torch.randn(3, 4)
print(a)
x = a.ge(0.5)
print(x)
print(a[x])
运行结果:
tensor([[ 0.1638, 0.9582, -0.2464, -0.8064],
[ 1.8385, -2.0180, 0.8382, 1.0563],
[ 0.1587, -1.6653, -0.2057, 0.1316]])
tensor([[0, 1, 0, 0],
[1, 0, 1, 1],
[0, 0, 0, 0]], dtype=torch.uint8)
tensor([0.9582, 1.8385, 0.8382, 1.0563])
take索引
take索引是基于目标Tensor的flatten形式下的,即摊平后的Tensor的索引。
import torch
a = torch.tensor([[3, 7, 2], [2, 8, 3]])
print(a)
print(torch.take(a, torch.tensor([0, 1, 5])))
运行结果:
tensor([[3, 7, 2],
[2, 8, 3]])
tensor([3, 7, 3])
转自:https://blog.csdn.net/qq8993174/article/details/90178555
https://blog.csdn.net/shu15121856/article/details/87810372