pytorch索引与切片

indexing

a=torch.randn(4,3,28,28)
a[0].shape
# torch.Size([3, 28, 28])

a[0,0].shape
# torch.Size([28, 28])

a[0,0,2,4]
# tensor(0.6574)

select first/last N

这个a可以看成一个图片:[batch,RBG,h,w]

a.shape
# torch.Size([4, 3, 28, 28])


a[:2].shape
# torch.Size([2, 3, 28, 28])
# 这里指的是后两张图片，这里可以是a[:2]，也可以是a[:2,:,:,:]，
# 如果要加上的话，需要加个:

a[:2,:1].shape
#torch.Size([2, 1, 28, 28])
#这个返回的是后两张图片的第一个RBG通道的数值

a[:2,1:,:,:].shape
#torch.Size([2, 2, 28, 28])
#这里返回的也是后两张图片的后两个RBG通道的数值

a[:2,-1:,:,:].shape
# torch.Size([2, 1, 28, 28])

这里需要注意两点：

• 1 就是这里的向前索引都是从0开始的并且不包括是前闭后开的

• 2.如果正着索引的话是比如[0,1,2]，倒是索引就是[-3,-2,-1]，所以这个a[:2,:-1,:,:].shape返回的就是[2,2,28,28]，因为不包括最后一个所以是2

select by steps

其实只有一种通用的方式就是\(begin:end:step\),这个代表着从\([begin,end)\)隔step取一个，而这种\(begin:end\)是另step=1的写法

a[:,:,0:28:2,0:28:2].shape
# torch.Size([4, 3, 14, 14])

a[:,:,::2,::2].shape
#torch.Size([4, 3, 14, 14])

select by specific index

这个就是我们可以选择不同的index,比如说选择第0张图片的第0，2个RBG

a.index_select(dim,torch.tensor([,,..])).shape
这个第一个dim是维度，然后就是在dim维度上选择[,,,...]列表中的索引

a.shape
#torch.Size([4, 3, 28, 28])

a.index_select(0,torch.tensor([0,2])).shape
#torch.Size([2, 3, 28, 28])
#这个就是在选择第0和第2张图片，idm就选择0


a.index_select(1,torch.tensor([0,1])).shape
#torch.Size([4, 2, 28, 28])
#这个就是选择全部图片的第0和第1个RBG,所以返回[4,2,28,28]

a.index_select(2,torch.arange(28)).shape
#torch.Size([4, 3, 28, 28])
#这个就是第二个维度上选择所有


a.index_select(2,torch.arange(8)).shape
#torch.Size([4, 3, 8, 28])
#这个就是在第二个维度上选择前8个

...(任意的的维度)

这个符号只是为了方便，比如说我们a[0,:,:,::2]，这里我们可以简写成a[0,...,::2],如果中间有很多冒号的话会很麻烦，所以可以用这个...来代替

a.shape
# torch.Size([4, 3, 28, 28])


a[...].shape
# torch.Size([4, 3, 28, 28])



a[0,...].shape
# torch.Size([3, 28, 28])



a[:,1,...].shape
# torch.Size([4, 28, 28])


a[...,:2].shape
# torch.Size([4, 3, 28, 2])



a[0,...,::2].shape
# torch.Size([3, 28, 14])

select by mask

这个用的比较少

.masked_select()
比如说torch.masked_select(x,x.ge(0.5))
这个就是返回x中数值大于0.5的元素