Different optimizers for shared parameters

username: newto
cooked: <p>Hi,</p>
<p>I’ve seen some similar questions but not exactly the same, and I haven’t found an answer.<br>
I have an encoder decoder model with some reconstruction loss, and the outputs of the encoder also go to another model and another loss –<br>
encoded = encoder(x)<br>
decoded = decoder(encoded)<br>
loss_recon = ||decoded - x||^2<br>
y = model2(encoded)<br>
loss_model = ||y-target||^2</p>
<p>Now, I want to update the encoder decoder with some optimizer (say optimizer 1 that has the encoder and decoder parameters) with respect to the reconstruction loss, and in addition, I want to update the encoder with the loss from the model, via a second optimizer (say optimizer 2 that has only the encoder parameters).<br>
Now, optimizer1 and 2 share the parameters of the encoder. How can I do the optimization step with a single pass through the graph?<br>
For instance, using<br>
optimizer1.zero_grad()<br>
loss_recon.backward(retain_graph=True)<br>
optimizer1.step()</p>
<p>optimizer2.zero_grad()<br>
loss_model.backward()<br>
optimizer2.step()</p>
<p>doesn’t work, because optimizer1.step() changed the parameters of the encoder.</p>
<p>using<br>
optimizer1.zero_grad()<br>
optimizer2.zero_grad()</p>
<p>loss_recon.backward(retain_graph=True)<br>
loss_model.backward()</p>
<p>optimizer1.step()<br>
optimizer2.step()</p>
<p>doesn’t work since only the grads from the second backward will remain.</p>
<p>Any ideas?</p>
<p>Edit:<br>
To simplify the question – assume I have a single model with two losses, and I want to use a different optimizer for each loss. Is there anyway I can to this with a single pass through the model?</p>
<p>optimzer1 = Adam(model.parameters(), lr=lr1)<br>
optimizer2 = Adam(model.parameters(), lr=lr2)<br>
out = model(x)<br>
loss1 = loss1(x, labels1) # update with respect to this loss with optimizer 1<br>
loss2 = loss2(x, labels2) # update with respect to this loss with optimizer 2</p>

username: ptrblck
cooked: <p>Your use case is a bit tricky and you’ve already posted some complications.<br>
Let’s check your summary:</p>
<pre><code class="lang-auto">optimzer1 = Adam(model.parameters(), lr=lr1)
optimizer2 = Adam(model.parameters(), lr=lr2)
out = model(x)
loss1 = loss1(x, labels1) # update with respect to this loss with optimizer 1
loss2 = loss2(x, labels2) # update with respect to this loss with optimizer 2
</code></pre>
<p>Your use case uses two different losses: <code>loss1</code> and <code>loss2</code>. These losses would calculate different gradients (<code>grad</code>) in the same model parameters (<code>param</code>) (assuming the forward activations (<code>act</code>) are still alive and valid).<br>
The optimizers should however use only the “corresponding” gradients for their update.</p>
<p>A potential approach could be to:</p>
<ul>
<li>calculate <code>grad1</code> from <code>loss1</code> using <code>act1</code> and <code>param1</code>
</li>
<li>store <code>grad1</code> separately and zero out the <code>param.grad</code> attributes</li>
<li>calculate <code>grad2</code> from <code>loss2</code> using <code>act1</code> and <code>param1</code>
</li>
<li>update the model using <code>opt2</code> to create parameter set <code>param2</code>
</li>
<li>zero out the gradients, restore the <code>.grad</code> attributes from the previously stored <code>grad1</code> tensors</li>
<li>update !!! <code>param2</code> !!! using <code>opt1</code> and <code>grad1</code>
</li>
</ul>
<p>Note that mathematically the last step is again wrong, since <code>grad1</code> was never calculated using these parameters.<br>
I don’t know what your exact use case is, but it seems that you want to update the model twice using stale gradients and losses, so could you explain the idea a bit more?</p>

username: newto
cooked: <p>Hi,</p>
<p>Thanks for the quick reply! Yes, I guess what you suggested will work while a bit hacky.</p>
<p>Maybe I’ll provide a different view for my use case, and there will be a better solution:<br>
The problematic point is that I have two losses to train an encoder, and one of the losses is used to train another model that uses the encoder’s output, but has a different scale to it –<br>
I have an encoder decoder model that uses some loss (say <code>loss1</code>). The encoded vectors are then used<br>
as an input to another model and there’s another loss at the output of that model (say <code>loss2</code>).<br>
Now, I want to update the encoder parameters using <code>loss1+0.001*loss2</code> and to update the parameters of the second model (the one that uses the encoder output, using <code>loss2</code>). The problem here is that there’s a different scale for <code>loss2</code> in the update of the encoder (<code>*0.001</code>), and the second model (<code>*1</code>).<br>
So, I can’t simply use 3 optimizers (encoder, decoder, model) and <code>loss=loss1+0.001*loss2</code>, since it will incorrectly update the model. Any good solution in this case? I can compensate for this by increasing the learning rate of model 2, but that is not equivalent, and seems to give bad results.<br>
Any ideas here?</p>

username: newto
cooked: <p>And another question, how would you store and load the gradients from the model efficiently?</p>

username: InnovArul
cooked: <p>I could think of the following solution. See if it works for your use case.</p>
<ol>
<li>First, do <code>loss1.backward(retain_graph=True)</code>. This will accumulate <code>loss1</code>’s gradients in encoder and decoder, and keep the computation graph intact.</li>
<li>Next, configure a backward hook on the encoded vectors to multiply <code>0.001</code> to its gradients.</li>
<li>Now, do <code>loss2.backward()</code>. Conceptually, this will backpropagate <code>loss2</code>’s gradients to second model and backpropagate <code>0.001 * loss2</code> gradients to the encoder.</li>
</ol>
<p>Hope it helps!</p>

username: newto
cooked: <p>Yes, this was one of the first things I tried, but the second backward (<code>loss2.backward()</code>) replaces the gradients from the first backward (and not accumulates them), so it ends up updating only with the gradients from <code>loss2</code>.</p>

username: InnovArul
cooked: <p>I am not sure why it didn’t work for you.<br>
Pytorch will not replace the gradients, but will aggregate the gradients. Perhaps you may have to check your code if you called a <code>zero_grad()</code> somewhere?</p>
<p>The following code shows the solution that I meant:</p>
<aside class="onebox githubgist" data-onebox-src="https://gist.github.com/InnovArul/845e909c770e9d15723590ed66f4d6ce">
  <header class="source">

      <a href="https://gist.github.com/InnovArul/845e909c770e9d15723590ed66f4d6ce" target="_blank" rel="noopener">gist.github.com</a>
  </header>

  <article class="onebox-body">
    <h4><a href="https://gist.github.com/InnovArul/845e909c770e9d15723590ed66f4d6ce" target="_blank" rel="noopener">https://gist.github.com/InnovArul/845e909c770e9d15723590ed66f4d6ce</a></h4>

  <h5>multi_loss_optimizer.py</h5>
  <pre><code class="Python">import torch, torch.nn as nn
import torch.optim as optim

def print_grads(modules, string):
    print(string)
    for mod in modules:
        for p in mod.parameters():
            print(p.grad)
        print('**')
    print("-----")</code></pre>
    This file has been truncated. <a href="https://gist.github.com/InnovArul/845e909c770e9d15723590ed66f4d6ce" target="_blank" rel="noopener">show original</a>

<p>
</p>

  </article>

  <div class="onebox-metadata">
    
    
  </div>

  <div style="clear: both"></div>
</aside>