PyTorch differentiable mask

Viewed 1460

How would I go about blacking out a portion of an image or feature map such that AutoGrad can backprop through the operation?

Specifically I want to black out everything except for n layers of border pixels. So if we consider a single channel of the feature map which looks like:

[
  [1, 1, 1, 1],
  [1, 1, 1, 1],
  [1, 1, 1, 1],
  [1, 1, 1, 1],
]

I set a constant n=1 so my operation does the following to the input:

[
  [1, 1, 1, 1],
  [1, 0, 0, 1],
  [1, 0, 0, 1],
  [1, 1, 1, 1],
]

In my case I'd be doing it to a multi channel feature map and all channels would be treated the same way.

If possible, I want to do it in a functional manner.

1 Answers

Considering the comments you added, i.e. that you don't need the output to be differentiable wrt. to the mask (said differently, the mask is constant), you could just store the indices of the 1s in the mask and act only on the corresponding elements of whatever Tensor you're considering. Or if you don't want to deal with fancy indexing, you could just keep the mask as a Tensor of 0s and 1s and do an element-wise multiplication of it with whatever Tensor you're considering. Or, if you truly just need to compute a loss along just the border pixels, just extract the first and last row, and first and last column, and avoid double-counting the corners. This latter solution is essentially just the first solution recast in a special case.

To address the question in your comment to my answer:

x = torch.tensor([[1.0,2,3],[4,5,6]], requires_grad = True)
print(x[:,0])

gives

tensor([1., 4.], grad_fn=<SelectBackward>)

, so we see that slicing does not mess with the autograd engine (it's still tracking the contribution to the gradient). It is not too surprising that this works automatically; slicing can be viewed as the (mathematical) function that of projecting onto a subspace of R^n, for which it's easy to compute the gradient.

Related