PyTorch_张量索引操作

简单行，列索引操作

import torch 
import numpy as np # 简单行列索引
def test01():data = torch.randint(0, 10, [4, 5])print(data)# 获得指定的某行元素print(data[0])# 获得指定某个列的元素print(data[:, 0])  # 逗号前面表示行，逗号后面表示列。冒号表示所有行或者所有列# 获得指定位置的某个元素print(data[1, 2])# 表示先获得前三行，再获得第三列的数据print(data[:3, 2])# 表示获得前三行的前两列print(data[:3, :2])# 列表索引
def test02():data = torch.randint(0, 10, [4, 5])print(data)# 如果索引的行列都是一个一维的列表，那么两个列表的长度必须相等# 表示获得 (0, 0), (2, 1), (3, 2) 三个位置的元素print(data[[0, 2, 3], [0, 1, 2]])# 表示获得 0, 2, 3 行的 0, 1, 2 列print(data[[[0], [2], [3]], [0, 1, 2]])if __name__ == "__main__":test02() 

布尔索引

import torch 
import numpy as np # 布尔索引
def test01():torch.manual_seed(0)data = torch.randint(0, 10, [4, 5])print(data)# 能够获得该张量中所有大于3的元素# 张量可以与数字做比较print(data > 3)print(data[data > 3])# 返回第2列元素大于6的行print(data[data[:, 1] > 6])# 返回第2行元素大于3的所有列print(data[:, data[1] > 3])# 多维索引
def test02():torch.manual_seed(0)data = torch.randint(0, 10, [3, 4, 5])print(data)# 选择第0行的所有元素print(data[0, :, :])# 按照第1哥维度选择第0元素, 是按行print(data[:, 0, :])# 按照第2个维度选择第0元素, 是按列print(data[:, :, 0])if __name__ == "__main__":test02()