pytorch的索引与切片

**

一、从左边开始索引

**

In [12]: a = torch.rand(4,3,28,28)

In [13]: a.shape
Out[13]: torch.Size([4, 3, 28, 28])

In [14]: a[0].shape
Out[14]: torch.Size([3, 28, 28])

In [15]: a[0,0].shape
Out[15]: torch.Size([28, 28])

In [16]: a[0,0,2,4]
Out[16]: tensor(0.4559)

**

二、选前或后的数据

**

In [17]: a.shape
Out[17]: torch.Size([4, 3, 28, 28])

In [18]: a[:2].shape#dim 0索引0,1
Out[18]: torch.Size([2, 3, 28, 28])

In [19]: a[:2,:1,:,:].shape#dim0索引0,1 dim1索引0
Out[19]: torch.Size([2, 1, 28, 28])

In [20]: a[:2,1:,:,:].shape#dim0索引0,1 dim1索引1到last
Out[20]: torch.Size([2, 2, 28, 28])

In [21]: a[:2,-1:,:,:].shape#dim0索引0,1 dim1索引-1到last
Out[21]: torch.Size([2, 1, 28, 28])

**

三、步进索引

**
start?step包含start，不包含end，步进为step

In [22]: a[:,:,::2,::2].shape#dim2步进2 dim3步进2
Out[22]: torch.Size([4, 3, 14, 14])

In [23]: a[:,:,0:28:2,0:28:2].shape#dim2步进2 dim3步进2
Out[23]: torch.Size([4, 3, 14, 14])

**

四、具体索引

**
a.index_select(0,torch.tensor([0,2]))第一个参数为对a的哪个维度进行操作，第二个参数是选取特定的索引

In [24]: a.shape
Out[24]: torch.Size([4, 3, 28, 28])

In [25]: a.index_select(0,torch.tensor([0,2])).shape#选取0维度的0,2
Out[25]: torch.Size([2, 3, 28, 28])

In [26]: a.index_select(1,torch.tensor([0,2])).shape#选取1维度的0,2
Out[26]: torch.Size([4, 2, 28, 28])

In [27]: a.index_select(2,torch.arange(28)).shape#选取2维度的前28
Out[27]: torch.Size([4, 3, 28, 28])

In [28]: a.index_select(2,torch.arange(8)).shape#选取2维度的前8
Out[28]: torch.Size([4, 3, 8, 28])

In [29]: a.shape
Out[29]: torch.Size([4, 3, 28, 28])

In [30]: a[...].shape
Out[30]: torch.Size([4, 3, 28, 28])

In [31]: a[0,...].shape
Out[31]: torch.Size([3, 28, 28])

In [32]: a[:,1,...].shape
Out[32]: torch.Size([4, 28, 28])

In [33]: a[...,:2].shape
Out[33]: torch.Size([4, 3, 28, 2])

**

五、掩码切片

**

In [34]: x = torch.randn(3,4)

In [35]: x
Out[35]:
tensor([[ 0.2171, -0.4370, -0.9547, -0.1153],
        [ 0.1009,  1.4227,  0.9283, -0.4541],
        [-0.8437, -0.7954, -1.8106, -1.8765]])

In [36]: mask = x.ge(0.5)#x中大于0.5的掩码

In [37]: mask
Out[37]:
tensor([[0, 0, 0, 0],
        [0, 1, 1, 0],
        [0, 0, 0, 0]], dtype=torch.uint8)

In [39]: torch.masked_select(x,mask)#掩码切片
Out[39]: tensor([1.4227, 0.9283])

In [40]: torch.masked_select(x,mask).shape
Out[40]: torch.Size([2])

**

六、打平索引

**

In [41]: src = torch.tensor([[4,3,5],[6,7,8]])

In [42]: torch.take(src,torch.tensor([0,2,5]))#先打平，再按打平索引
Out[42]: tensor([4, 5, 8])