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

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.


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.



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

Ok, thank you for confirming.

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)

Will backpropagation still work after using such a masking operation?

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