目录
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])))
|