pytorch：索引与切片

参考网易云课堂的pytorch 学习

#indexing dim 0 first 
import torch
a=torch.rand(4,3,28,28)
a[0].shape#第0张的图片的c h w
#torch.Size([3, 28, 28])
a[0,0].shape #第0张图片的 0通道的 h,w
a[0,0,2,4]#第0张图片的 0通道 的第[2,4]的像素值
a[:2,:1,:,:].shape #前两个图片 从0 开始左闭右开 [0,2) [0,1) [0,28) [0,28)
a[:2,-1:,:,:].shape # [0,2) 最后一个元素取到最末尾
#select by specific index 在对应的维度选择
a.shape
a.index_select(0,torch.tensor([0,2])).shape #在第0维 即batch 中选择 选择第0个和第2个图片
# torch.Size([2, 3, 28, 28])
a.index_select(1,torch.tensor([1,2])).shape #在第1维选择第1个通道和第2个通道
# torch.Size([4, 2, 28, 28])
a.index_select(2,torch.arange(28)).shape #在第2个维度选择[0,28) 即全部选择
# torch.Size([4, 3, 28, 28])
a.index_select(2,torch.arange(8)).shape #在第2个维度选在[0,8)
# torch.Size([4, 3, 8, 28])
#...
a.shape
# torch.Size([4, 3, 28, 28])
a[0,...].shape #从第0 维选择第一个，即选择第一张图片
# torch.Size([3, 28, 28])
a[:,1,...].shape #选择第一个通道的所有图片
# torch.Size([4, 28, 28])
a[...,:2].shape #在最后一维选择[0,2)
# torch.Size([4, 3, 28, 2])
#select by mask .masked_select
x=torch.randn(3,4)
x
# tensor([[-0.0775, -0.6049,  0.5763, -0.1832],
#         [-1.1918,  0.7328,  0.3736, -0.0185],
#         [-0.2479, -1.2790,  0.1330, -0.4590]])

mask=x.ge(0.5) #大于0.5的设置对应矩阵的位置设置为True
mask #bool 类型
# tensor([[False, False,  True, False],
#         [False,  True, False, False],
#         [False, False, False, False]])

torch.masked_select(x,mask) #最后在相应为True的元素的值取出来设置成1维向量
# tensor([0.5763, 0.7328])