6.7 张量索引操作

6.7.1 简单索引

import torch

tensor1 = torch.randint(1, 9, (3, 5, 4))
print(tensor1)
print()

# 取 第0维第0
print(tensor1[0])
print()

# 取 第0维所有，第1维第1
print(tensor1[:, 1])
print()

# 取 第0维所有，第1维第1，第2维第3
print(tensor1[2, 1, 3])

6.7.2 范围索引

import torch

tensor1 = torch.randint(1, 9, (3, 5, 4))
print(tensor1)
print()

# 取 第0维第1到最后
print(tensor1[1:])
print()

# 取 第0维最后，第1维1到3(包含3),第2维0到2(包含2)
print(tensor1[-1:, 1:4, 0:3])
print()

6.7.3 列表索引

import torch

tensor1 = torch.randint(1, 9, (3, 5, 4))
print(tensor1)
print()

# 取 第0维第0，第1维第1 和 第0维第1，第1维第2
print(tensor1[[0, 1], [1, 2]])
print()

# 取 第0维第0，第1维第1、2 和 第0维第1，第1维第1、2
print(tensor1[[[0], [1]], [1, 2]])

6.7.4 布尔索引

import torch

tensor1 = torch.randint(1, 9, (3, 5, 4))
print(tensor1)
print()

# 取 第2维第0大于5的，返回(dim0,dim1)形状的索引
print(tensor1[:, :, 0] > 5)
print(tensor1[tensor1[:, :, 0] > 5])
print()

# 取 第1维第1大于5的，返回(dim0,dim2)形状的索引
mask = tensor1[:, 1, :] > 5
print(mask)
tensor2 = tensor1.permute(0, 2, 1)  # 转换维度为(dim0,dim2,dim1)
print(tensor2[mask])
tensor2 = tensor2[mask].permute(1, 0)  # 转换维度为(dim1,?)
print(tensor2)
print()

# 取 第1维第1，第2维第2大于5的，返回(dim0)形状的索引
print(tensor1[:, 1, 2] > 5)
print(tensor1[tensor1[:, 1, 2] > 5])