Compute gradient wrt different weights

username: jojono
cooked: <p>How can I compute the gradient w.r.t to weights different from the one used to compute the loss.</p>
<p>To clarify my problem: I have a neural network model and its weights W1. I train my model on some data and so the weights changed and are now W2. Then I compute the loss L using W2 on a new set of data. Now I would like to update my weight W1 using gradient computed w.r.t to W1 on L (which was computed using W2).</p>
<p>I tried to reload the weights W1 using load_state_dict between the compute of the loss and loss.backward() but it does not work. I think that the problem is that load_state_dict change some internal states and variables in addition to the weights.</p>
<p>How can I solve my problem ?<br>
Thank you</p>

username: albanD
cooked: <p>Hi,</p>
<p>Keep in mind that <code>state_dict()</code> does not copy your weights. So after you update your model to contain W2, your state dict also contains W2.<br>
You might want to do it by hand with <code>saved_params = [p.clone() for p in model.parameters()]</code> and load it with</p>
<pre><code class="lang-auto">with torch.no_grad():
  for p, old_val in zip(model.parameters(), saved_params):
    p.copy_(old_val)
</code></pre>

username: jojono
cooked: <p>Thank you or the explanation of state_dict.</p>
<p>However, I found this thread <a href="https://discuss.pytorch.org/t/understanding-load-state-dict-effect-on-computational-graph/16115/6" class="inline-onebox">Understanding load_state_dict() effect on computational graph</a> where they explained that my solution is not the good one. But I do not understand how to use nn.functional interface to change the parameters w.r.t which the gradients will be computed during the backward pass.<br>
Can you help me ?</p>

username: albanD
cooked: <p>The custom reloading of the weights above should do the right thing for you.<br>
You don’t want to perform forward with custom weights here, you want to snapshot and reload the weights right?</p>

username: jojono
cooked: <p>I want to compute the gradient of a loss computed using W2, w.r.t W1. Where W2 is the current state of the network and W1 a previous state.</p>

username: albanD
cooked: <p>I’m not sure to understand then.<br>
Do you want the gradient computation to take into account the first forward/backward, the optimizer step and the second forward?<br>
Or you just want to use gradients computes with W2 and apply them to W1?</p>

username: jojono
cooked: <p>Sorry, it was not clear.</p>
<p>W1 and W2 are both weights of the neural network but at two distinct moment.<br>
L is the loss function<br>
grad(W, x) is the gradient of x w.r.t. W<br>
net(W, x) is the output of the neural network for x as input when the network is in state W.</p>
<p>I want to compute grad(W1, L(net(W2,x))) where x corresponds to training data.</p>

username: albanD
cooked: <p>How do you define W2 here?<br>
Do you have <code>W2 = grad_step(W1, grad(W1, L(net(W1, x)))</code> or <code>W2 = grad_step(W1, grad(W1, L(net(W1, x))).detach()</code>?</p>

username: jojono
cooked: <p>W2 is defined as W2 = grad_step(W1, grad(W1, L(net(W1, x))).</p>

username: albanD
cooked: <p>So this is going to be a bit more tricky.<br>
The first backward with W1 will need to set <code>create_graph=True</code>.<br>
But then you won’t be able to use optimizer to do the step as they operate inplace and without registering their operation to the autograd.<br>
So you either need to reimplement the optimizer to do this operation in a differentiable manner. And assign these new values without deleting the old ones and keeping gradients flowing (which can be tricky).<br>
Or check libraries such as <a href="https://github.com/facebookresearch/higher" rel="nofollow noopener"><code>higher</code></a> that already do similar things <img src="https://discuss.pytorch.org/images/emoji/apple/smiley.png?v=9" title=":smiley:" class="emoji" alt=":smiley:"></p>