Backpropagate through gradients generated by autograd

username: Eugene_Lee
cooked: <p>Hi there. Let’s say we obtain a set of gradients for a set of weights as follows:</p>
<pre><code class="lang-auto">grad = torch.autograd.grad([output], [m for m in self.parameters()],
                                    grad_outputs=[grad_input],
                                    create_graph=True, retain_graph=retain_graph,
                                    only_inputs=True)
for i, m in enumerate(self.parameters()):
        m.data = m.data - lr * grad[i]
</code></pre>
<p>Let the model that had its weight updated by <code>net</code>. We then perform a forward pass using the weight updated model as <code>net(input)</code>. As I’ve retained the graph in the <code>autograd</code> performed above, when I’m performing <code>loss.backward()</code>, I’d like the gradient to flow through the <code>grad_input</code> (backproped from the <code>grad[i]</code>'s) variable mentioned above. Any idea how to achieve this?</p>
<p>Here’s a simple flow of the program:</p>
<pre><code class="lang-auto">gradOut = self.netA(input_data)
self.netB.paramUpdate(gradOut) # performs the code snippet provided above
out = self.netB(input_data)
loss = self.criterion(out, label)
self.optimizer.zero_grad()
loss.backward() # I hope that the gradients can flow from netB to netA
</code></pre>

username: albanD
cooked: <p>HI,</p>
<p>Hi, the problem is that you use <code>.data</code> which breaks that history and prevent gradient computation.</p>
<p>In particular, if netB does not have any learnable parameters, you should remove them (del netB.some_mod.weight for example) and then fill it in with the value computed: netB.some_mod.weight = lr * grad[i]</p>

username: Eugene_Lee
cooked: <p>Hi <a class="mention" href="/u/alband">@albanD</a>, thanks for the reply.</p>
<p>The reason of using <code>.data</code> is because if I don’t use it, I realize the the parameter of <code>netB</code> won’t be updated.</p>
<p>Both <code>netA</code> and <code>netB</code> have learnable parameters. The idea is that when I’m doing a forward pass on <code>netB</code>, its parameter is a composition of <code>netB</code>'s original parameter plus a gradient provided by <code>netA</code>. So when I’m doing backprop, I’m hoping that the gradient will update both <code>netB</code> and <code>netA</code> as the provided gradient should be part of the computational graph.</p>

username: Eugene_Lee
cooked: <p>I found that <a href="https://github.com/tristandeleu/pytorch-meta" rel="nofollow noopener">pytorch-meta</a> has a similar approach where they pre-compute the updated parameter and re-insert it into the existing model. Is this the most elegant way to accomplish this?</p>

username: albanD
cooked: <p>You will have to differentiate between the leant parameter in netB and the ones that are computed at each forward.<br>
One thing you cal do is rename parameters like weight into weight_orig. Then del the weight. And in the forward part, set the weight as weight_orig + your_udpate.</p>

username: Eugene_Lee
cooked: <p><a class="mention" href="/u/alband">@albanD</a> Thanks! This seems like the only way to get this done. It’s working fine now using this implementation.</p>