Given an array and mask of same shapes, I want the masked output of the same shape and containing 0 where mask is False.
For example,
# input array img = torch.randn(2, 2) print(img) # tensor([[0.4684, 0.8316], # [0.8635, 0.4228]]) print(img.shape) # torch.Size([2, 2]) # mask mask = torch.BoolTensor(2, 2) print(mask) # tensor([[False, True], # [ True, True]]) print(mask.shape) # torch.Size([2, 2]) # expected masked output of shape 2x2 # tensor([[0, 0.8316], # [0.8635, 0.4228]]) Issue: The masking changes the shape of the output as follows:
#1: shape changed img[mask] # tensor([0.8316, 0.8635, 0.4228]) 3 Answers
Simply type-cast your boolean mask to an integer mask, followed by float to bring the mask to the same type as in img. Perform element-wise multiplication afterwards.
masked_output = img * mask.int().float()
One of the ways I found to solve it was:
img[mask==False] = 0 or using
img[~mask] = 0 It'll change the img itself.
The most straight forward way would be creating another tensor to handle it.
import torch def generate_masked_tensor(input, mask, fill=0): masked_tensor = torch.zeros(input.size()) + fill masked_tensor[mask] = input[mask] return masked_tensor if __name__ == "__main__": img = torch.randn(2, 2) mask = torch.tensor([False, True, True, False]).bool().view(2, 2) masked_img = generate_masked_tensor(img, mask) print (masked_img) The output:
tensor([[0.0000, 0.8028], [1.5411, 0.0000]]) 
