5分鐘玩轉(zhuǎn)PyTorch | PyTorch張量的索引方法
↑?關(guān)注 + 星標?，每天學Python新技能
后臺回復【大禮包】送你Python自學大禮包
在掌握常用符號索引的基礎(chǔ)上，最好可以適應(yīng)函數(shù)式索引的思路。因為Pytorch中很多函數(shù)都采用的是函數(shù)式索引的思路，而且使用函數(shù)式索引對代碼可讀性會有很大提升。本文看似基礎(chǔ)，其實都是細節(jié)。
1 張量的符號索引
張量也是有序序列，我們可以根據(jù)每個元素在系統(tǒng)內(nèi)的順序位置，來找出特定的元素，也就是索引。
1.1 一維張量的索引
一維張量由零維張量構(gòu)成
一維張量索引與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)
注：張量索引出的結(jié)果是零維張量，而不是單獨的數(shù)。要轉(zhuǎn)化成單獨的數(shù)還需使用上節(jié)介紹的item()方法。
可理解為構(gòu)成一維張量的是零維張量，而不是單獨的數(shù)。
張量的step必須大于0
#?索引3-10號元素，左閉右開，默認step為1
t1[2:?8]
#?tensor([3,?4,?5,?6,?7,?8])
#?step=3，隔3個數(shù)取一個，左閉右開
t1[2:?8:?2]
#?tensor([3,?5,?7])
在Python中，step可以為負數(shù)，例如:
li?=?[1,?2,?3]
#?列表倒敘排列，取所有數(shù)值，從后往前取
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]的索引結(jié)果就不同：
t2[::2,?::2]
#?tensor([[?1,??3],
#????????[?9,?11]])
t2[::2][::2]
#?tensor([[1,?2,?3,?4]])
t2[::2, ::2]二維索引使用逗號隔開時，可以理解為全局索引，取第一行和第三行的第一列和第三列的元素。
t2[::2][::2]二維索引在兩個中括號中時，可以理解為先取了第一行和第三行，構(gòu)成一個新的二維張量，然后在此基礎(chǔ)上又間隔2并對所有張量進行索引。
tt?=?t2[::2]
#?tensor([[?1,??2,??3,??4],
#?????????[?9,?10,?11,?12]])
tt[::2]
#?tensor([[1,?2,?3,?4]])
1.3 三維張量的索引
設(shè)三維張量的shape是x、y、z，則可理解為它是由x個二維張量構(gòu)成，每個二維張量由y個一維張量構(gòu)成，每個一維張量由z個元素構(gòu)成。
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 張量的函數(shù)索引
2.1 一維張量的函數(shù)索引
在PyTorch中，我們還可以使用index_select函數(shù)指定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這個一維向量來說，由于只有一個維度，第二個參數(shù)取值為0，就代表在第一個維度上進行索引，索引的位置是1和2。
注:這里取出的是位置，而不是取出[1:2]區(qū)間內(nèi)左閉右開的元素。
2.2 二維張量的函數(shù)索引
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參數(shù)取值為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參數(shù)取值為1，代表在shape的第二個維度上進行索引。index參數(shù)的值為[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次。
高維張量函數(shù)索引的思路與低維一樣，都是在shape的維度上進行操作。
在PyTorch中很多函數(shù)都采用的是第幾維的思路，后面會介紹給大家，大家還需勤加練習，適應(yīng)這種思路。同時使用函數(shù)式索引，在習慣后對代碼可讀性會有很大提升。
推薦閱讀