Pytorch根据mask对张量进行筛选

最新推荐文章于 2022-04-29 16:14:05 发布

飞机火车巴雷特

最新推荐文章于 2022-04-29 16:14:05 发布

阅读量2.2k

点赞数 2

分类专栏： Python相关

本文链接：https://blog.csdn.net/qq_36158230/article/details/121666048

版权

pytorch

Python相关专栏收录该内容

61 篇文章 1 订阅

订阅专栏

一、需求

输入

[-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.])