Updating 2 models with shared loss

username: mEVPj1KVN4MM
cooked: <p>I have 2 models that share 1 part of the loss, I have to minimize one loss and maximize other one.</p>
<pre><code class="lang-auto">opt1 = optimizer(m1.parameters())
opt2 = optimizer(m2.parameters())

for epoch in range(epochs):
    opt1.zero_grad()
    opt2.zero_grad()

    m1 = model1(x)
    m2 = model2(x)

    loss1 = shared_loss(m1, m2) + loss_fn1(m1)
    loss2 = -1*(shared_loss(m1, m2) + loss_fn2(m2))  # -1 because loss2 should be maximized
</code></pre>
<p>How would I implement the backward pass in this case?<br>
I can’t just do</p>
<pre><code class="lang-auto">loss1.backward()
loss2.backward()
opt1.step()
opt2.step()
</code></pre>
<p>because that would go through the graph twice.</p>

username: KFrank
cooked: <p>Hi m!</p>
<aside class="quote no-group" data-username="mEVPj1KVN4MM" data-post="1" data-topic="165559" 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/m/ea666f/48.png" class="avatar"> mEVPj1KVN4MM:</div>
<blockquote>
<pre><code class="lang-auto">    loss1 = shared_loss(m1, m2) + loss_fn1(m1)
    loss2 = -1*(shared_loss(m1, m2) + loss_fn2(m2))  # -1 because loss2 should be maximized
</code></pre>
<p>How would I implement the backward pass in this case?</p>
</blockquote>
</aside>
<p>It’s not entirely clear to me what you want to do here.</p>
<p>I assume that you want to optimize the parameters of <code>model1</code> so as to<br>
minimize <code>loss1</code>, thus, in part, minimizing <code>shared_loss</code>, and optimize<br>
the parameters of <code>model2</code> so as to minimize <code>loss2</code>, which includes<br>
<code>-shared_loss</code>, therefore implying maximizing <code>shared_loss</code>.</p>
<p>That is, leaving aside <code>loss_fn1</code> and <code>loss_fn2</code> as they don’t change<br>
the basic issue, my understanding is that you want <code>model1</code> to minimize<br>
<code>shared_loss</code> while <code>model2</code> maximizes <code>shared_loss</code>.</p>
<p>If so, the cleanest approach will be to use a single backward pass and<br>
a single optimizer, but to flip the gradient of <code>shared_loss</code> when you<br>
backpropagate through <code>model2</code>.  So something like:</p>
<pre data-code-wrap="python"><code class="lang-python">opt = torch.optim.SGD (list (model1.parameters()) + list (model2.parameters()), lr = learning_rate)
...
m1 = model1 (x)
m2 = model2 (x)
loss = shared_loss (m1, GradientReversalFunction.apply (m2))
opt.zero_grad()
loss.backward()
opt.step()
</code></pre>
<p>Here <code>GradientReversalFunction</code> is a custom autograd function that passes<br>
the value of <code>m2</code> through unchanged during the forward pass, but changes<br>
the sign of the gradient of <code>m2</code> during the backward pass.</p>
<p>This thread discusses a couple of ways to effect such a gradient reversal:</p>
<aside class="quote" data-post="1" data-topic="134810">
  <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/hadaev8/48/16280_2.png" class="avatar">
    <a href="https://discuss.pytorch.org/t/how-to-reverse-gradient-sign-during-backprop/134810">How to reverse gradient sign during backprop?</a> 
  </div>
  <blockquote>
    For example, I have models net0 and net1 working like this 
out0 = net0(inp0)
out1 = net0(net1(net0(inp1)))
loss = criterion(out0, out1, target)

I want to reverse the gradient sign from the output of model net1 but calculate net0(inp0) as usual. 
In simple case, I would do 
out0 = net0(inp0)
out1 = net0(net1(inp1))
loss = criterion(out0, out1, target)
loss.backward()
[p.grad.data.neg_() for p in net1.parameters()]
opt.step()

Can I do some hook to reverse gradients from output of net1?
  </blockquote>
</aside>

<p>Best.</p>
<p>K. Frank</p>