Autograd.grad breaks in hypernetwork implementations

username: Arkaprava_Majumdar
cooked: <p>Hi guys, I am trying to make a hypernetwork and got stuck.</p>
<p>z0 = torch.Tensor([1.,1.])<br>
input_size = z0.shape[0]</p>
<p>mainNetwork = torch.nn.Sequential(<br>
torch.nn.Linear(input_size,64),<br>
torch.nn.Tanh(),<br>
torch.nn.Linear(64,input_size))</p>
<p>p_shape,_,theta_init = get_parameters(mainNetwork)<br>
input_size_theta = theta_init.shape[0]</p>
<p>hyperNetwork = torch.nn.Sequential(<br>
torch.nn.Linear(input_size_theta,64),<br>
torch.nn.Tanh(),<br>
torch.nn.Linear(64,input_size_theta))<br>
z0 = torch.tensor([1.,1.],requires_grad = True)<br>
t = 0<br>
theta_0 = hyperNetwork(theta_init)<br>
set_parameters(mainNetwork,theta_0)<br>
z1 = mainNetwork(z0)<br>
torch.autograd.grad(z1,hyperNetwork.parameters(),grad_outputs = torch.ones_like(z1),allow_unused = True)<br>
I get (None,None,None,None) as output.<br>
get_parameters is a function that returns all the flattened parameters of a model into a single tensor.<br>
set_parameters takes a single tensor and sets it as the model weights.</p>
<p>def get_parameters(model):<br>
p_shape = <span class="chcklst-box checked fa fa-check-square-o fa-fw"></span><br>
flat_parameters = <span class="chcklst-box checked fa fa-check-square-o fa-fw"></span><br>
theta = torch.empty(0)<br>
for p in model.parameters():<br>
p_shape.append(p.size())<br>
flat_parameters.append(p.flatten())<br>
theta = torch.cat([theta,p.flatten()],dim = 0)<br>
return p_shape,flat_parameters,theta</p>
<p>def set_parameters(model,theta):<br>
p_shape,flat_parameters,_ = get_parameters(model)<br>
idx = 0<br>
j = 0<br>
for i in range(len(model)):<br>
if isinstance(model[i],torch.nn.Linear):<br>
sub_theta_weight = theta[idx : idx + np.prod(p_shape[j])].reshape(p_shape[j])<br>
sub_theta_bias = theta[idx + np.prod(p_shape[j]):idx + np.prod(p_shape[j]) + np.prod(p_shape[j+1])].reshape(p_shape[j+1])<br>
model[i].weight = torch.nn.Parameter(sub_theta_weight)<br>
model[i].bias = torch.nn.Parameter(sub_theta_bias)<br>
idx = np.prod(p_shape[j]) + np.prod(p_shape[j+1])<br>
j += 2</p>
<p>I also went through <a href="https://discuss.pytorch.org/t/hypernetwork-implementation/152109/8" class="inline-onebox">Hypernetwork implementation - #8 by mariaalfaroc</a><br>
Apparently, even in these examples, torch.autograd.grad(x,hyperNetwork.parameters()) is giving the exact same error as in my case even though the loss.backward() is working fine.<br>
Any help will be greatly appreciated. I am new to autograd and trying to learn more about how it works.<br>
Thanks.</p>

username: KFrank
cooked: <p>Hi Arkaprava!</p>
<aside class="quote no-group quote-modified" data-username="Arkaprava_Majumdar" data-post="1" data-topic="187898" data-full="true">
<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/arkaprava_majumdar/48/61353_2.png" class="avatar"> Arkaprava_Majumdar:</div>
<blockquote>
<p>theta_0 = hyperNetwork(theta_init)<br>
set_parameters(mainNetwork,theta_0)<br>
z1 = mainNetwork(z0)<br>
torch.autograd.grad(z1,hyperNetwork.parameters(),grad_outputs = torch.ones_like(z1),allow_unused = True)<br>
I get (None,None,None,None) as output.<br>
…<br>
model[i].weight = torch.nn.Parameter(sub_theta_weight)<br>
model[i].bias = torch.nn.Parameter(sub_theta_bias)</p>
</blockquote>
</aside>
<p>I don’t follow your code in detail, but it appears that <code>theta_0</code> depends on<br>
the parameters of <code>hyperNetwork</code>, that the values of the parameters of<br>
<code>mainNetwork</code> depend on the value of <code>theta_0</code>, and that <code>z1</code> depends on<br>
<code>mainNetwork</code>.</p>
<p>However, because of how <code>set_parameters()</code> is implemented, the computation<br>
graph does not connect the parameters of <code>mainNetwork</code> to <code>theta_0</code>.  So the<br>
computation graph that you might imagine connects <code>z1</code> to <code>hyperNetwork</code> is<br>
“broken.”  Therefore <code>autograd.grad (z1, hyperNetwork.parameters())</code><br>
gives you <code>None</code>s.</p>
<p>I don’t see this documented anywhere, but, although a new <code>Parameter</code> that is<br>
constructed from an existing <code>Parameter</code> shares the original <code>Parameter</code>’s data,<br>
it is not the same <code>Tensor</code> nor does it share the original <code>Parameter</code>’s <code>grad</code>.<br>
Furthermore, autograd does not track the connection between the new and<br>
original <code>Parameter</code>s.  That is, constructing a new <code>Parameter</code> from an existing<br>
<code>Parameter</code> “breaks the computation graph.”</p>
<p>Consider:</p>
<pre><code class="lang-plaintext">&gt;&gt;&gt; import torch
&gt;&gt;&gt; torch.__version__
'2.0.1'
&gt;&gt;&gt; p0 = torch.nn.Parameter (torch.ones (3))
&gt;&gt;&gt; p1 = torch.nn.Parameter (p0)
&gt;&gt;&gt; p1.sum().backward()
&gt;&gt;&gt; p1.grad
tensor([1., 1., 1.])
&gt;&gt;&gt; p0.grad
&gt;&gt;&gt;
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

username: Arkaprava_Majumdar
cooked: <p>Thanks. Is there any way to get around this?</p>

username: KFrank
cooked: <p>Hi Arkaprava!</p>
<aside class="quote no-group" data-username="Arkaprava_Majumdar" data-post="3" data-topic="187898" data-full="true">
<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/arkaprava_majumdar/48/61353_2.png" class="avatar"> Arkaprava_Majumdar:</div>
<blockquote>
<p>Is there any way to get around this?</p>
</blockquote>
</aside>
<p>Probably – but it depends on what you’re actually trying to accomplish.</p>
<p>The problem is that <code>Parameter</code>s don’t like to be modified while being tracked<br>
by autograd.</p>
<p>Here is an example script this addresses this core issue:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

class MainNetwork (torch.nn.Module):
    def __init__(self, input_size):
        super().__init__()
        # no Parameters nor requires_grad = True
        self.lin1_weight = torch.zeros (64, input_size)
        self.lin1_bias = torch.zeros (64)
        self.lin2_weight = torch.zeros (input_size, 64)
        self.lin2_bias = torch.zeros (input_size)
    
    def forward (self, x):
        # use functional.linear() instead of Linear
        x = torch.nn.functional.linear (x, self.lin1_weight, self.lin1_bias)
        x = torch.tanh (x)
        x = torch.nn.functional.linear (x, self.lin2_weight, self.lin2_bias)
        return x

def set_parameters (model, w1, b1, w2, b2):
    # don't recreate any Parameters -- just assign Tensors
    model.lin1_weight = w1
    model.lin1_bias = b1
    model.lin2_weight = w2
    model.lin2_bias = b2

input_size = 2
z0 = torch.ones (input_size, requires_grad = True)
w1 = torch.ones (64, input_size, requires_grad = True)
b1 = torch.ones (64, requires_grad = True)
w2 = torch.ones (input_size, 64, requires_grad = True)
b2 = torch.ones (input_size, requires_grad = True)

mainNetwork = MainNetwork (input_size)
set_parameters (mainNetwork, w1, b1, w2, b2)

z1 = mainNetwork (z0)
print (z0, z1)

w1_grad, b1_grad, w2_grad, b2_grad = torch.autograd.grad (z1, (w1, b1, w2, b2), grad_outputs = torch.ones_like(z1), allow_unused = True)

print (w1_grad.shape, b1_grad.shape, w2_grad.shape, b2_grad.shape)
</code></pre>
<p>And here is its output:</p>
<pre><code class="lang-plaintext">2.0.1
tensor([1., 1.], requires_grad=True) tensor([64.6835, 64.6835], grad_fn=&lt;AddBackward0&gt;)
torch.Size([64, 2]) torch.Size([64]) torch.Size([2, 64]) torch.Size([2])
</code></pre>
<p>If this approach doesn’t work for your use case, please post a fully-self-contained,<br>
runnable script that illustrates a simplified version of your issue, together with the<br>
output you get when you run it.</p>
<p>Best.</p>
<p>K. Frank</p>

username: vaibhav
cooked: <p>Hi KFrank,</p>
<p>thank you for this answer. any way to use this functional API for RNN? have to create RNN from scratch?</p>