How to fill in a batch of tensor with another tensor according to mask

vision
1

Hi all,
Suppose I have the following code

to_learn = nn.Parameter(torch.randn([1, 3, 10, 10]))

In the forward, I have a mask where each mask masks different areas of a tensor


def forward(x, mask)
    bs = x.shape[0]
    new= torch.zeros([bs, 3, 64, 64]).cuda()
    # fill in new with to_learn where mask is 1
    ...
    return x + new

The shapes of mask and x are [bs, 3, 64, 64].
How can I efficiently fill in the tensor new with the to_learn tensor only at locations where mask==1?

Thank you for any help.

2

You don’t need to allocate a new tensor for this as a separate step, as an out of place operation should achieve the same thing, e.g.,

new = torch.mul(to_learn, mask)

A similar approach is used by the pruning utilities:
https://pytorch.org/tutorials/intermediate/pruning_tutorial.html

3

Hi thank you for the suggestion. But this will cause size not matching issue.
RuntimeError: The size of tensor a (64) must match the size of tensor b (10) at non-singleton dimension 3
Do you happen to know another way?

4

Ah, I interpreted the question as intending to mask to_learn, but if the mask is intended for the input, you can do

x = torch.mul(x, mask)
5

Thank you. There was a typo in my original post I have fixed.

My purpose is to fill in a batch of image x with to_learn when mask is 1. Suppose there is an image of a dog. The mask gives me a segment of the dog’s head. I want to fill in the head area with to_learn.

Using the suggested method will not fill in x with to_learn right?

6

No, it achieves the goal of masking the input preventing the parameters in to_learn from accessing the masked values.

7

Thank you. In this case, x is masked. Then how do we fill in to_learn into masked area of x considering x and to_learn in different shape?

8

If you mean how to predict the output with a masked input x, then this is doable with a common operation like a convolution.