pytorch中mask操作之torch.masked_select

# Copyright 2018 CVTE . All Rights Reserved.
# coding: utf-8

from __future__ import print_function
from __future__ import division
from __future__ import absolute_import
import torch

# torch.masked_select(input, mask, *, out=None) → Tensor
# 返回input中mask为True的元素，组成1维tensor
# input 和mask的维度不一定要相同,但是需要可以扩展成同样维度

x = torch.randn(3, 4)
print(x)
"""
tensor([[-0.4152, -1.1952,  0.5442,  1.3843],
        [ 0.5764,  0.6684, -0.6330,  0.4646],
        [ 1.1162, -0.8730,  0.0285,  0.7387]])
"""

mask = x.ge(0.5)
print(mask)

"""
tensor([[0, 0, 1, 1],
        [1, 1, 0, 0],
        [1, 0, 0, 1]], dtype=torch.uint8)
"""

print(torch.masked_select(x, mask))

"""
tensor([0.5442, 1.3843, 0.5764, 0.6684, 1.1162, 0.7387])
"""

# input 和mask的维度不一定要相同,但是需要可以扩展成同样维度
print(torch.masked_select(x, torch.tensor([False, False, True, False])))
"""
tensor([ 0.5442, -0.6330,  0.0285])
"""
print(torch.masked_select(x, torch.tensor([[False, False, True, False]])))
"""
tensor([ 0.5442, -0.6330,  0.0285])
"""
print(torch.masked_select(x, torch.tensor([[False], [False], [True]])))
"""
tensor([ 1.1162, -0.8730,  0.0285,  0.7387])
"""