Masking out locations in convolutional kernels

username: mjdmahsneh
cooked: <p>Hello Pytorchers!</p>
<p>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.</p>
<p>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).</p>
<p><strong>Now my question comes in two parts:</strong></p>
<ol>
<li>
<p>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).</p>
</li>
<li>
<p>Is there a way to also customize the bias term, such that it is only added to locations of interest?</p>
</li>
</ol>
<pre><code class="lang-auto">class MaskedConv3d(nn.Module):
        def __init__(self, n_in, n_out, filter_mask, pad = 1):
            super().__init__()
            self.kernel_size = tuple(filter_mask.shape)
            self.register_buffer('filter_mask', filter_mask)
            self.conv = nn.Conv3d(
                in_channels=n_in,
                out_channels=n_out,
                kernel_size=self.kernel_size,
                stride=1,
                bias=False,
                padding = pad
            )
        def forward(self, x):
            self._mask_conv_filter()
            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)

</code></pre>

username: eqy
cooked: <p>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.</p>

username: mjdmahsneh
cooked: <p><a class="mention" href="/u/eqy">@eqy</a> 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!</p>
<p>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!</p>

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

username: mjdmahsneh
cooked: <p>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.</p>

username: eqy
cooked: <p>Ok, do you mean that you want no bias to be applied for input patches that look like</p>
<pre><code class="lang-auto">1 0 1
0 0 0
1 0 1
</code></pre>
<p>? 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.</p>

username: mjdmahsneh
cooked: <p>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.</p>

username: Vasquez122
cooked: <p>This is a great inspiring article. I am pretty much pleased with your good work<a href="https://www.upsers.run/" rel="noopener nofollow ugc">!</a></p>

username: mjdmahsneh
cooked: <p>Thank you <a class="mention" href="/u/vasquez122">@Vasquez122</a>, I actually built upon <a href="https://discuss.pytorch.org/t/applying-custom-mask-on-kernel-for-cnn/87099">Applying custom mask on kernel for CNN</a>. Thanks to their efforts.</p>

username: Ahmed_Tarek
cooked: <p>Thank you for sharing your implementation, it was very helpful</p>
<p>I have a minor suggestion that can simplify the code. We can actually skip this section:</p>
<pre><code class="lang-auto">with torch.no_grad():
                self.conv.weight = nn.Parameter(self.conv.weight * self.filter_mask)```
</code></pre>
<p>My understanding is that <code>register_buffer</code> treats the tensor as a non-trainable parameter, so wrapping it with <code>nn.Parameter</code> and using <code>torch.no_grad</code> isn’t really necessary.</p>
<p>Here’s how the code could look after simplifying it:</p>
<pre><code class="lang-auto">class MaskedConv2d(nn.Module):
    def __init__(self,cin,cout,kernel, mask,padding='same'):
        super().__init__()
        self.conv_ = nn.Conv2d(in_channels = cin,
                                out_channels = cout,
                                kernel_size = kernel,
                                padding=padding,
                               bias=False)
        self.register_buffer('mask', mask)
    def forward(self,x):
        self.conv_.weight.data = self.conv_.weight.data * self.mask
        return self.conv_(x)
</code></pre>