Computation graph is interrupted when calling a model within a model

gchou: <p>Hi all,</p>
<p>I am trying to do the following:</p>
<ol>
<li>Create a model (model A) with a few linear layers that produces X, a list of tensors</li>
<li>Create a model (model B) with a few linear layers during the forward loop of model A (or create model B one time outside of the training loop and just reference it here each time)</li>
<li>Overwrite the parameters of model B with X</li>
</ol>
<pre><code class="lang-auto">idx = 0
for n, p in model_B.named_parameters():
    w = X[idx].reshape(128, 128)
    idx += 1 
    p.data = w.float() 
</code></pre>
<ol start="4">
<li>Run model B on some input Y given the new parameters</li>
</ol>
<pre><code class="lang-auto">out = model_B(Y)
loss = F.l1_loss(out, ground_truth)
</code></pre>
<ol start="5">
<li>Finally, call <code>loss.backward()</code> in the forward loop of model A and update the parameters of model A</li>
</ol>
<p>Currently, the gradient of the parameters of model A is None. Maybe the graph is interrupted when I try to replace the parameters of model B with X as <code>p.data = w.float()</code> is implicitly performed with <code>no_grad()</code>? I’m not sure if this is the issue though.</p>
<p>Some additional information:</p>
<ol>
<li>X[0]…X[t] all have requires_grad == True</li>
<li>all of model_B parameters have is_leaf == True and requires_grad == True</li>
<li>if I comment out the code for model B and run something simple such as loss = X[0].mean(), then the gradients of the parameters of model A are valid, so the problem probably isn’t with model A itself</li>
</ol>
<p>Any advice would be greatly appreciated, thanks!</p>

KFrank: <p>Hi gchou!</p>
<aside class="quote no-group" data-username="gchou" data-post="1" data-topic="158898" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/g/91b2a8/48.png" class="avatar"> gchou:</div>
<blockquote>
<p>Maybe the graph is interrupted when I try to replace the parameters of model B with X</p>
</blockquote>
</aside>
<p>Models typically have <code>Parameter</code>s and <code>Parameter</code>s are typically things<br>
that are optimized.</p>
<p>If I understand your use case, you don’t intend to optimize the “parameters”<br>
(weights) of model B, but rather derive (“predict”) them using another model,<br>
model A (whose <code>Parameter</code>s you do intend to optimize).</p>
<p>There are a lot of ways of packaging this, but I think that writing model B<br>
as a function, rather than structuring is as an “official model” (i.e., <code>Module</code>),<br>
would be the cleanest approach and would communicate more clearly what<br>
you are doing.</p>
<p>The following abbreviated pseudo-code illustrates this approach:</p>
<pre data-code-wrap="python"><code class="lang-python">def functional_model_b (input, weights):   # takes an input tensor and a list of four weight tensors
    t = torch.nn.functional.linear (input, weight = weights[0], bias = weights[1]
    t = torch.nn.functiona.relu (t)
    t = torch.nn.functional.linear (t, weight = weights[2], bias = weights[3])
    return t

model_b_weights = model_A (some_input)   # returns list of four tensors
pred = functional_model_b (some_other_input, model_b_weights)
loss = criterion (pred, targ)
loss.backward()   # backpropagates through functional_model_b and computes gradients for model_A's parameters
</code></pre>
<p>Best.</p>
<p>K. Frank</p>

gchou: <p>Hi K. Frank,</p>
<p>Thanks so much for the response, this would work perfectly for what I need to do!<br>
Much cleaner than what I was going for as well <img src="https://discuss.pytorch.org/images/emoji/apple/slight_smile.png?v=12" title=":slight_smile:" class="emoji" alt=":slight_smile:" loading="lazy" width="20" height="20"></p>