pytorch深度学习入门—tensor的索引与数据筛选

目录

torch.where

torch.index_select

torch.gather

torch.masked_select

torch.take

---

---

torch.where

参数含义：满足这个condition则用这个tensor的元素，否则则用另一个

代码如下：

# torch.where
a = torch.rand(4, 4)
b = torch.rand(4, 4)
print(a, b)
# 满足这个condition则用这个tensor的元素，否则则用另一个
print(torch.where((a > 0.5), a, b))

---

torch.index_select

代码如下：

# torch.index_select
print("------------torch.index_select-------------")
print(a)
print(torch.index_select(a, dim=1, index=torch.tensor([0, 3, 2])))

---

torch.gather

dim维度解释：第0维是列，第1维是行，第2维是行中元素展开……

代码如下：

# torch.gather
print("==========torch.gather==========")
a = torch.linspace(1, 16, 16).view(4, 4)
print(a)
# 第0维是列，第1维是行，第2维是行中元素展开……
print(torch.gather(a, dim=0, index=torch.tensor([[0, 1, 1, 1], [0, 1, 2, 2], ])))

torch.masked_select

参数解释：第一个参数为要筛选的tensor，第二个参数为boolen Tensor，将true的值输出为一个向量

 代码如下：

# torch.masked_select
print("------------torch.masked_index-----------")
mask = torch.gt(a, 8)
print(a)
print("mask:", mask)
# 第一个参数为要筛选的tensor，第二个参数为boolen Tensor，将true的值输出为一个向量
print(torch.masked_select(a, mask))

---

torch.take

代码如下：

# torch.take
print("-------------------")
print(torch.take(a, index=torch.tensor([0, 1, 14, 3])))