5分钟玩转PyTorch | PyTorch张量的索引方法
共 3391字，需浏览 7分钟
·
2021-12-02 12:11
↑ 关注 + 星标 ，每天学Python新技能
后台回复【大礼包】送你Python自学大礼包
在掌握常用符号索引的基础上，最好可以适应函数式索引的思路。因为Pytorch中很多函数都采用的是函数式索引的思路，而且使用函数式索引对代码可读性会有很大提升。本文看似基础，其实都是细节。
1 张量的符号索引
张量也是有序序列，我们可以根据每个元素在系统内的顺序位置，来找出特定的元素，也就是索引。
1.1 一维张量的索引
一维张量由零维张量构成
一维张量索引与Python中的索引一样是是从左到右，从0开始的，遵循格式为[start: end: step]。
t1 = torch.arange(1, 11)
t1
# tensor([ 1, 2, 3, 4, 5, 6, 7, 8, 9, 10])
# 取出索引位置是0的元素
t1[0]
# tensor(1)
注：张量索引出的结果是零维张量，而不是单独的数。要转化成单独的数还需使用上节介绍的item()方法。
可理解为构成一维张量的是零维张量，而不是单独的数。
张量的step必须大于0
# 索引3-10号元素，左闭右开，默认step为1
t1[2: 8]
# tensor([3, 4, 5, 6, 7, 8])
# step=3，隔3个数取一个，左闭右开
t1[2: 8: 2]
# tensor([3, 5, 7])
在Python中，step可以为负数，例如:
li = [1, 2, 3]
# 列表倒叙排列，取所有数值，从后往前取
li[ ::-1]
# [3, 2, 1]
但在张量中，step必须大于1，否则就会报错。
t1 = torch.arange(1, 11)
t1[ ::-1]
# ValueError: step must be greater than zero
1.2 二维张量的索引
二维张量的索引逻辑和一维张量的索引逻辑相同，二维张量可以视为两个一维张量组合而成。
t2 = torch.arange(1, 17).reshape(4, 4)
t2
#tensor([[ 1, 2, 3, 4],
# [ 5, 6, 7, 8],
# [ 9, 10, 11, 12],
# [13, 14, 15, 16]])
t2[0,1]也可用t2[0][1]的表示。
# 表示索引第一行、第二个（第二列的）元素
t2[0, 1]
# tensor(2)
t2[0][1]
# tensor(2)
但是t2[::2, ::2]与t2[::2][ ::2]的索引结果就不同：
t2[::2, ::2]
# tensor([[ 1, 3],
# [ 9, 11]])
t2[::2][::2]
# tensor([[1, 2, 3, 4]])
t2[::2, ::2]二维索引使用逗号隔开时，可以理解为全局索引，取第一行和第三行的第一列和第三列的元素。
t2[::2][::2]二维索引在两个中括号中时，可以理解为先取了第一行和第三行，构成一个新的二维张量，然后在此基础上又间隔2并对所有张量进行索引。
tt = t2[::2]
# tensor([[ 1, 2, 3, 4],
# [ 9, 10, 11, 12]])
tt[::2]
# tensor([[1, 2, 3, 4]])
1.3 三维张量的索引
设三维张量的shape是x、y、z，则可理解为它是由x个二维张量构成，每个二维张量由y个一维张量构成，每个一维张量由z个元素构成。
t3 = torch.arange(1, 28).reshape(3, 3, 3)
t3
# tensor([[[ 1, 2, 3],
# [ 4, 5, 6],
# [ 7, 8, 9]],
# [[10, 11, 12],
# [13, 14, 15],
# [16, 17, 18]],
# [[19, 20, 21],
# [22, 23, 24],
# [25, 26, 27]]])
# 索引第二个矩阵中的第二行、第二个元素
t3[1, 1, 1]
# tensor(14)
# 索引第二个矩阵，行和列都是每隔两个取一个
t3[1, ::2, ::2]
# tensor([[10, 12],
# [16, 18]])
高维张量的思路与低维一样，就是围绕张量的“形状”进行索引。
2 张量的函数索引
2.1 一维张量的函数索引
在PyTorch中，我们还可以使用index_select函数指定index来对张量进行索引，index的类型必须为Tensor。
index_select(dim, index)表示在张量的哪个维度进行索引，索引的位值是多少。
t1 = torch.arange(1, 11)
indices = torch.tensor([1, 2])
# tensor([1, 2])
t1.index_select(0, indices)
# tensor([2, 3])
对于t1这个一维向量来说，由于只有一个维度，第二个参数取值为0，就代表在第一个维度上进行索引，索引的位置是1和2。
注:这里取出的是位置，而不是取出[1:2]区间内左闭右开的元素。
2.2 二维张量的函数索引
t2 = torch.arange(12).reshape(4, 3)
t2
# tensor([[ 0, 1, 2],
# [ 3, 4, 5],
# [ 6, 7, 8],
# [ 9, 10, 11]])
t2.shape
# torch.Size([4, 3])
indices = torch.tensor([1, 2])
t2.index_select(0,indices)
# tensor([[3, 4, 5],
# [6, 7, 8]])
此时dim参数取值为0，代表在shape的第一个维度上进行索引。
t2 = torch.arange(12).reshape(4, 3)
indices = torch.tensor([1, 1])
t2.index_select(1, indices)
# tensor([[ 1, 1],
# [ 4, 4],
# [ 7, 7],
# [10, 10]])
此时dim参数取值为1，代表在shape的第二个维度上进行索引。index参数的值为[1,1]，就代表取出第二个维度上为1的元素2次。
下面可以再次理解：
t2 = torch.arange(12).reshape(4, 3)
t2
# tensor([[ 0, 1, 2],
# [ 3, 4, 5],
# [ 6, 7, 8],
# [ 9, 10, 11]])
t2.shape
# torch.Size([4, 3])
indices = torch.tensor([2, 2, 2])
t2.index_select(1, indices)
# tensor([[ 2, 2, 2],
# [ 5, 5, 5],
# [ 8, 8, 8],
# [11, 11, 11]])
取出第二个维度上为2的元素3次。
高维张量函数索引的思路与低维一样，都是在shape的维度上进行操作。
在PyTorch中很多函数都采用的是第几维的思路，后面会介绍给大家，大家还需勤加练习，适应这种思路。同时使用函数式索引，在习惯后对代码可读性会有很大提升。
推荐阅读