Equivalent of numpy.ma.array to mask values in pytorch?

1

Does anyone know how to quickly apply a mask to a torch tensor? Say I have a tensor t_in with some values in it and a mask t_mask with values 0 and 1. In numpy I could do the following:

t_result = np.ma.array(t_in, mask=t_mask)

That would mask all the values of t_in where the value of the equivalent index of t_mask is 0. Then I can do operations such as t_result.sum() and it will only give me the sum of the non-masked values. I’m primarily looking to do just that, but cannot get my head around how it can be done efficiently with torch tensors.

Thanks

3

Hmm ok so after digging more looks like something like this could work:

t_m_bool = t_mask.type(torch.ByteTensor)
t_result = t_in.masked_select(t_m_bool)

Would love it if someone else could confirm that this is the best approach.

Thanks

4

Hi,

If you just want to set the values to 0 when the mask is 0 then the mask and your tensor should be of the same type and just multiply the two.
If you want to extract as a new tensor all the entries of your tensor where the mask value is 1, the the mask should be a ByteTensor and you can use masked_select or []: tensor[mask] or tensor.masked_select(mask).

5

Ok, thank you for confirming.

6

I recently found the need for doing masked operations on PyTorch tensors. Here’s how I implemented them (taken from my blog post):

def masked_mean(tensor, mask, dim):
    """Finding the mean along dim"""
    masked = torch.mul(tensor, mask)  # Apply the mask using an element-wise multiply
    return masked.sum(dim=dim) / mask.sum(dim=dim)  # Find the average!

def masked_max(tensor, mask, dim):
    """Finding the max along dim"""
    masked = torch.mul(tensor, mask)
    neg_inf = torch.zeros_like(tensor)
    neg_inf[~mask] = -math.inf  # Place the smallest values possible in masked positions
    return (masked + neg_inf).max(dim=dim)
7

Will backpropagation still work after using such a masking operation?

8

Yes it will.
The values that are used will get gradients corresponding to how they’re used and the values that are zeroed out will get 0 gradient.

1 Like