[人工智能]Pytorch根据mask对张量进行筛选

一、需求

? ? ? ? 输入

[-0.9, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0]

? ? ? ? 期望输出(得到一个仅包含正数或其他想要的元素的张量)

[1.0, 2.0, 3.0, 4.0]

二、实现

? ? ? ? 代码

import torch

# input
a = torch.tensor([-0.9, 0.0, 0.0, 1.0, 2.0, 3.0, 4.0])
print(a)
print()

# get the mask
non_zero = (a != 0)
less_zero = (a < 0)
zero = (a == 0)
larger_zero = (a > 0)

# output: select by mask
print(non_zero, torch.masked_select(a, non_zero))
print(less_zero, torch.masked_select(a, less_zero))
print(zero, torch.masked_select(a, zero))
print(larger_zero, torch.masked_select(a, larger_zero))  # 获得仅包含正值元素的张量

? ? ? ? 输出

tensor([-0.9000,  0.0000,  0.0000,  1.0000,  2.0000,  3.0000,  4.0000])

tensor([ True, False, False,  True,  True,  True,  True]) tensor([-0.9000,  1.0000,  2.0000,  3.0000,  4.0000])
tensor([ True, False, False, False, False, False, False]) tensor([-0.9000])
tensor([False,  True,  True, False, False, False, False]) tensor([0., 0.])
tensor([False, False, False,  True,  True,  True,  True]) tensor([1., 2., 3., 4.])