Masking out locations in convolutional kernels

Hello Pytorchers!

I am trying to implement a 3D convolutional layer where kernels have some sampling locations completely masked out. Particularly, I want to pass a binary mask, such that locations that are set to zero do NOT contribute to the learning process.

In the example below, I am using a cross-shaped mask, then multiplying it by the convolutional kernel weights, to eliminate responses from zeroed positions (where mask == 0).

Now my question comes in two parts:

  1. Am I achieving this correctly? Mind that my aim is to keep using this constant mask during all forward/backward passes, with no updates on the mask weights (so that at any point during training, only locations of interest are being learned).

  2. Is there a way to also customize the bias term, such that it is only added to locations of interest?

class MaskedConv3d(nn.Module):
        def __init__(self, n_in, n_out, filter_mask, pad = 1):
            self.kernel_size = tuple(filter_mask.shape)
            self.register_buffer('filter_mask', filter_mask)
            self.conv = nn.Conv3d(
                padding = pad
        def forward(self, x):
            return self.conv(x)
        def _mask_conv_filter(self):
            with torch.no_grad():
                self.conv.weight = nn.Parameter(self.conv.weight * self.filter_mask)

# define mask:
mask = torch.tensor([[[0., 1., 0.],
                      [1., 1., 1.],
                      [0., 1., 0.]],

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

                     [[0., 1., 0.],
                      [1., 1., 1.],
                      [0., 1., 0.]]])
x = torch.randn(1, 1, 6, 6, 6)

maskedconv3d = MaskedConv3d(1, 8, mask)
out = maskedconv3d(x)

I don’t see any obvious problems here but you can do some simple tests like running your layer on a ones tensor input and checking that the results are what you expect based on the mask. If you are using batchnorm layers after the convolution, you can avoid the bias term entirely as it will be effectively undone by the batchnorm. Additionally, I don’t think the bias is applied before the convolution, so it shouldn’t be affected by (or affect) the mask that you are using.

@eqy Thank you for your response. As you suggested, I have run a few tests to make sure am getting the expected output, it seems to be working as expected!

However, the fact that the bias is added after the convolution process, means that positions where mask(convolution(x))==0, will evaluate to 0+bias (after adding the bias term). So for now, I have set bias=False but it would be nice to also customize the bias term along with the masking process!

I’m not sure I understand what you mean by positions where mask(convolution(x))==0, as in the current code looks like the mask is independent of position (the same cross pattern is applied everywhere).

Sorry, it wasn’t the best way to explain my point. I meant to say that the response of the convolutional kernel would have values that are equal to zero (the masked out locations). These will have the bias term added to them and hence will evaluate to 0+bias. Kindly correct me if am missing something.

Ok, do you mean that you want no bias to be applied for input patches that look like

1 0 1
0 0 0
1 0 1

? I think that is trickier to achieve since it almost looks like some kind of data-dependent control flow. However, I would be surprised if this made a large difference in the accuracy of a model though, as these are exactly the kind of patterns that I would expect a convolutional model to learn around/adapt to.

I agree, I think it would have a minimal effect on the model. Thank you for your response and help, it is very much appreciated.

This is a great inspiring article. I am pretty much pleased with your good work!

1 Like

Thank you @Vasquez122, I actually built upon Applying custom mask on kernel for CNN. Thanks to their efforts.