pytorch学习05：索引和切片

索引

import torch

a = torch.rand(4,3,28, 28)
print("a.shape: ", a.shape)
print("a[0].shape: ", a[0].shape)
print("a[0][0].shape: ", a[0][0].shape)
print("a[0, 0].shape: ", a[0, 0].shape)
print("a[0, 1, 2, 3]: ", a[0, 1, 2, 3])

取前N个或后N个

import torch

a = torch.rand(4,3,28, 28)
print("a.shape: ", a.shape)
# 取第一维[0, 2)
print("a[:2].shape: ", a[:2].shape)
# 取第一维的 [0, 2) ，第二维的 [0, 1)，第三维的全部，第四维的全部
print("a[:2, :1, :, :].shape: ", a[:2, :1, :, :].shape)
# 取第一维的 [0, 2) ，第二维 [1, 最后一个] ，第三维的全部，第四维的全部
print("a[:2, 1:, :, :].shape: ", a[:2, 1:, :, :].shape)
# 取第一维的 [0, 2) ，第二维 [倒数第1个, 最后一个] ，第三维的全部，第四维的全部
print("a[:2, -1:, :, :].shape: ", a[:2, -1:, :, :].shape)

按步长选择

import torch

a = torch.rand(4,3,28, 28)
# 第三维取[0, 20),步长为2, 第四维取[0, 20),步长为2
print("a[:, :, 0:20:2, 0:20:2].shape: ", a[:, :, 0:20:2, 0:20:2].shape)
# 对第三维和第四维所有数据按步长为2进行抽取
print("a[:, :, ::2, ::2].shape: ", a[:, :, ::2, ::2].shape)

b = torch.arange(0, 10)
print("b[::2]", b[::2])

自定义选取

import torch

a = torch.rand(4,3,28, 28)

# 第一个变量代表维度， 第二个变量代表索引
# 第一维只保留索引为1和2的数据
print("a.index_select(0, torch.tensor([1,2])).shape: ",
      a.index_select(0, torch.tensor([1,2])).shape)

# 第二维只保留索引为1的数据
print("a.index_select(1, torch.tensor([1])).shape: ",
      a.index_select(1, torch.tensor([1])).shape)

”…“符号

import torch

a = torch.rand(4,3,28, 28)
# ... 表示任意多维度，只能放在最前面或最后面
# 取所有维度
print("a[...].shape: ", a[...].shape)
# 取第一维第0个，其余维度全取，等价于a[0, :, :, :]
print("a[0, ...].shape: ", a[0, ...].shape)
# 取第一维全部，第二位第1个，其余为全取, 等价于a[:, 1, :, :]
print("a[:, 1, ...].shape: ", a[:, 1, ...].shape)
# 取最后一维前两个，其余维全取，等价于a[:, :, :, 2]
print("a[..., :2].shape: ", a[..., :2].shape)

mask选取

import torch

a = torch.rand(3, 4)
print("a:\n", a)

# 判断每个元素是否大于等于0.5
mask = a.ge(0.5)
print("mask:\n", mask)

# mask会打平数据
print("torch.masked_select(a, mask): ", torch.masked_select(a, mask))
print("torch.masked_select(a, mask).shape: ", torch.masked_select(a, mask).shape)

take选取

import torch

a = torch.rand(3, 4)
print("a:\n", a)

# 会先打平再选取， [3, 4] -> [12]
print("torch.take(a, torch.tensor([0, 2]):"
      , torch.take(a, torch.tensor([0, 2, 11])))