Masked nn.Linear

Gabriel_CRETIN: <p>I try to build my own class of <code>nn.Linear</code> called <code>MaskedLinear</code> that is supposed to set some weights to 0 and keep the others. The input is a matrix of shape (1024, 1024) and the masks is the same size matrix full of ones and zeros (the first n columns of the matrix = 1 and the last columns are full of 0).</p>
<p>Here is my implementation so far:</p>
<pre><code class="lang-python">class MaskedLinear(nn.Linear):
    def __init__(self, in_features, out_features, bias=True):
        super().__init__(in_features, out_features, bias)
        
    def forward(self, input, masks):
        # `input` is of shape (batch_size, 1024, 1024)
        masks = masks.view(-1, 1, 1024*1024)
        # Expand the tensor to enable multiplication with self.weight
        masks = masks.expand(-1, self.weight.size(0), 1024*1024)
        print("input ", input.shape)
        print("masks ", masks.shape)
        print("self.weight ", self.weight.shape)
        return F.linear(input, masks * self.weight, self.bias)
</code></pre>
<p>And here is the error I get:</p>
<pre><code class="lang-python">Cell In [36], line 14, in MaskedLinear.forward(self, input, masks)
     12 print("masks ", masks.shape)
     13 print("self.weight ", self.weight.shape)
---&gt; 14 return F.linear(input, masks * self.weight, self.bias)

RuntimeError: t() expects a tensor with &lt;= 2 dimensions, but self is 3D
</code></pre>
<p>The corresponding layer is:</p>
<pre><code class="lang-auto">MaskedLinear(1024*1024, 128)
</code></pre>
<p>When I print the shapes of <code>input</code>, <code>masks</code> and self.weight it looks like:</p>
<pre><code class="lang-auto">input  torch.Size([8, 1048576])
masks  torch.Size([8, 128, 1048576])
self.weight  torch.Size([128, 1048576])
</code></pre>
<p>First, I don’t understand why the batch dimension is not managed.<br>
Second question : is it the right way to mask the weights ?</p>

tom: <p>So do you want to mask the weight or the input?<br>
So the weight is necessarily 2d in F.linear, but you could use <code>matmul</code> or <code>einsum</code> or whatever fits your tastes (plus adding the bias) to replace linear.</p>
<p>Best regards</p>
<p>Thomas</p>

Gabriel_CRETIN: <p>I don’t want the network to learn the padded sections of the matrix (c.f the screenshot of 2 partial examples below: true images are up, below are the predicted ones. Dark blue represents padding, so is “false” masking).<br>
I use a custom MaskedMSE loss function, where I can pass my masking so that the network does not backpropagate errors on padded sections. The missing thing is to apply the masking in the layers of the network (only nn.Linear).</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/e/feaa82e20df9bffc54b2cf720c9b395dd4577d02.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/feaa82e20df9bffc54b2cf720c9b395dd4577d02" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/feaa82e20df9bffc54b2cf720c9b395dd4577d02_2_690x70.jpeg" alt="image" data-base62-sha1="AkSAXKdpFJGxV9agW7HCSQkJojM" width="690" height="70" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/feaa82e20df9bffc54b2cf720c9b395dd4577d02_2_690x70.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/feaa82e20df9bffc54b2cf720c9b395dd4577d02_2_1035x105.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/e/feaa82e20df9bffc54b2cf720c9b395dd4577d02_2_1380x140.jpeg 2x" data-dominant-color="4A488E"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">image</span><span class="informations">1822×185 60.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>I am not very familiar with <code>matmul</code> and <code>einsum</code>…</p>

tom: <p>I cannot comment on the design, but so if the input (not the weight) is masked to 0 in the forward pass, the gradient contribution to the weight that is applied to that masked part should should be 0 (because the backward computation matrix multiplies the gradient of the output with the input (in the right order, with the right transpositions) to get the weight gradient?<br>
If you have NaNs/infs that won’t work, though and you need to deal with those.</p>
<p>Best regards</p>
<p>Thomas</p>

Gabriel_CRETIN: <p>Sorry for the late reply.<br>
Indeed, I finally am able to mask the input in the forward pass so it is not necessary to do anything particular at the loss calculation for the reason you explained.<br>
Thanks for the explanations !</p>