Gradient not computed for nn.Parameter

Hiba_Ahsan: <p>I have a model with a linear layer and I wish to multiply the linear layer weights element-wise with a torch.nn.Parameter at every forward pass. I tried it the the following way but this not compute a gradient for the parameter.</p>
<p>class LinearModel(nn.Module):</p>
<pre><code>def __init__(self, args):
      super(LinearModel, self).__init__()
      self.linear1 = nn.Linear(10, 20)
      self.mask = torch.nn.Parameter(torch.randn_like(self.linear1.weight))

def forward(self, x)
      self.linear1.weight = self.linear1.weight *self.mask
      return self.linear1(x)
</code></pre>
<p>The loss is squared loss. When I print self.mask.grad, it shows None. Where am I going wrong?</p>

ptrblck: <p>Your code should not even be executable and fails with:</p>
<pre><code class="lang-python">lin = LinearModel(0)
x = torch.randn(1, 10)
out = lin(x)
# TypeError: cannot assign 'torch.FloatTensor' as parameter 'weight' (torch.nn.Parameter or None expected)
</code></pre>
<p>since you are trying to assign a tensor to a parameter.<br>
In case you want to manipulate the <code>linear1.weight</code> parameters inplace before executing the forward pass, wrap it into a <code>torch.no_grad()</code> guard or use the functional API via:</p>
<pre><code class="lang-python">class LinearModel(nn.Module):
    def __init__(self):
          super(LinearModel, self).__init__()
          self.linear1 = nn.Linear(10, 20)
          self.mask = torch.nn.Parameter(torch.randn_like(self.linear1.weight))
    
    def forward(self, x):
          weight = self.linear1.weight *self.mask
          return F.linear(x, weight, self.linear1.bias)


lin = LinearModel()
x = torch.randn(1, 10)
out = lin(x)
out.mean().backward()
print(lin.linear1.weight.grad)
print(lin.linear1.bias.grad)
print(lin.mask.grad)
</code></pre>

Hiba_Ahsan: <aside class="quote no-group" data-username="ptrblck" data-post="2" data-topic="169917">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/48/1823_2.png" class="avatar"> ptrblck:</div>
<blockquote>
<p>torch.no_grad()</p>
</blockquote>
</aside>
<p>Thank you! Sorry, I pasted an outdated code (I tried <code>self.linear1.weight = nn.Parameter(self.linear1.weight*self.mask</code>) but that gave None as the gradient for <code>mask</code>.</p>
<p>Can you please show how I can do the same using <code>torch.no_grad()</code> as you earlier suggested?</p>

srishti-git1110: <aside class="quote no-group" data-username="Hiba_Ahsan" data-post="3" data-topic="169917">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/hiba_ahsan/48/12120_2.png" class="avatar"> Hiba_Ahsan:</div>
<blockquote>
<p>Thank you! Sorry, I pasted an outdated code (I tried <code>self.linear1.weight = nn.Parameter(self.linear1.weight*self.mask</code>) but that gave None as the gradient for <code>mask</code>.</p>
</blockquote>
</aside>
<p>Hi,<br>
That is expected.<br>
As you are essentially creating a new parameter with this line of code and that would be a leaf tensor (with its <code>grad_fn=None</code>).<br>
This means <code>mask</code> is no more present in the computation graph of <code>linear1.weight</code> or <code>out</code> and hence does not get its <code>grad</code> populated.</p>
<p>See the following code - I have used torchviz to visualize the comp graphs :</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import torch.nn as nn

class LinearModel(nn.Module):
  def __init__(self):
      super(LinearModel, self).__init__()
      self.linear1 = nn.Linear(10, 20)
      self.mask = torch.nn.Parameter(torch.randn_like(self.linear1.weight))

  def forward(self, x):
      self.linear1.weight = nn.Parameter(self.linear1.weight*self.mask)
      print(self.linear1.weight.is_leaf) # True
      print(self.linear1.weight.grad_fn) # None
      make_dot(self.linear1.weight).render("weight", format="png") # image 1
      return self.linear1(x)


lin = LinearModel()
x = torch.randn(1, 10)
out = lin(x)
make_dot(out).render("output", format="png") # image 2
out.mean().backward()

print(lin.linear1.weight.grad)
print(lin.linear1.bias.grad)
print(lin.mask.grad) # None
</code></pre>
<p>Now, if you see image 1 below it is a leaf with no graph associated with it. Hence, mask isn’t present in both graphs anymore.</p>
<p>image 1 (linear1.weight)</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/9/e/9e7b511beb347500d94b4148030de57e324d95c6.png" alt="image" data-base62-sha1="mBZE10YVveVwxSA45xnzi8WDTDg" width="105" height="52"></p>
<p>image 2 (out)<br>
The top two blue nodes correspond to the bias and the weight term of the linear1 layer.</p>
<p><img src="https://discuss.pytorch.org/uploads/default/original/3X/e/8/e87c6ba24f489edc7b6858cc859f536b4d20c7a4.png" alt="image" data-base62-sha1="xaFlfpxJsjoO1Dz43e4halzx54g" width="288" height="377"></p>

Hiba_Ahsan: <p>Thank you! Is there a way I can do this without creating a new module? I don’t know how to do this with <code>torch.no_grad()</code> as suggested earlier. I do want the gradient to be computed for <code>mask</code>.</p>

srishti-git1110: <p>Is there any specific reason not to use the code posted by <a class="mention" href="/u/ptrblck">@ptrblck</a> above?</p>
<p>With torch.no_grad, you could definitely modify the weight of the linear1 layer by replacing <code>self.linear1.weight = nn.Parameter(self.linear1.weight*self.mask)</code> with -</p>
<pre><code class="lang-auto">with torch.no_grad():
     self.linear1.weight *= self.mask
</code></pre>
<p>But again, since autograd is looking away mask will not be included in the graph - grad will be None again.<br>
Not sure if there’s a way to modify linear1.weight in-place while also including mask in the graph.</p>

Hiba_Ahsan: <p>Makes sense. I’ll go ahead with what <a class="mention" href="/u/ptrblck">@ptrblck</a> suggested then. Thank you so much again!</p>