PyTorch Batch masked select implementation

user_2940718 0
Viewed 1039

How can I perform a batched masked_select?

Given:

x = torch.tensor([[1., 2., 2., 2., 3.],
                  [1., 2., 4., 3., 2.]])

The desired output would be:

tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])

Here is a possible approach:

x = torch.tensor([[1., 2., 2., 2., 3.],
                  [1., 2., 4., 3., 2.]])

ones = torch.tensor([[1., 1., 1., 1., 1.],
                     [1., 1., 1., 1., 1.]])

masks = torch.tensor([[ True, False, False, False,  True],
                      [ True, False,  True,  True, False]])

for i in range(x.size(0)):
    mask = masks[i]
    s = torch.masked_select(x[i], mask)
    ones[i][:s.size(0)] = s

Are there any alternative solutions?

pytorch-dataloader pytorch deep-learning
1 Answers

The main issue with those kinds of problems is that the intermediate result is non-homogenous: in your batch elements will have different numbers of masked values. This is an issue if we want to apply PyTorch builtins. Here I am providing two solutions to perform this operation.

1- Using a list comprehension

Go through the batch elements, mask and pad by the appropriate amount:

>>> pad = lambda v: F.pad(v, [0, len(m)-len(v)], value=1)
>>> torch.stack([pad(r[m]) for r, m in zip(x, masks)])
tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])

This is rather straightforward and similar to your approach.

2- Using torch.scatter

A vectorized alternative is to construct the correct value and index tensors such that we can apply torch.scatter and obtain the desired result. The trick here is to work with flattened tensors. From x and masks we first want to get access to nz and idx defined as:

  • nz: the non-masked values from x (using masks of course), i.e. we need to find:

    tensor([1., 3., 1., 4., 3.])

  • idx: their corresponding indices in the output tensor when flattened.

    tensor([ 0,  1,  5,  6,  7])

Then we can apply scatter with something like out = ones.scatter(dim=0, idx, nz) which will effectively perform: out[idx[i]] = nz[i].

To construct nz, we can directly index x with masksindices of non-zero values using on masks:

>>> nz = x[masks]
tensor([1., 3., 1., 4., 3.])

For idx it will be a little bit more tricky. We can sort the mask itself, flatten it and get the non-zero values using torch.Tensor.nonzero. When sorted, the True values end up at the beginning of each rows:

>>> idx = masks.sort(1, True).values.view(-1).nonzero()[:,0]
tensor([ 0,  1,  5,  6,  7])

Finally we can apply torch.scatter and reshape to get the desired result:

>>> torch.ones(x.numel()).scatter(0, idx, nz).view_as(x)
tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])

The usage of torch.scatter here is limited since the inputs are one-dimensional. An equivalent approach is to simply to:

>>> o = torch.ones(x.numel())
>>> o[idx] = nz
>>> o.view_as(x)

Complete approach:

>>> idx = masks.sort(1, True)[0].view(-1).nonzero()[:,0]
>>> torch.ones(x.numel()).scatter(0, idx, x[masks]).view_as(x)
tensor([[1., 3., 1., 1., 1.],
        [1., 4., 3., 1., 1.]])
user_63313690
Related