How to optimize model by combining multi-task losses

username: Max.T
cooked: <p>Hi, I am new to pytorch and not quite sure if I am doing right. Please tell me whether the following is right or wrong.</p>
<p>What I am doing:<br>
I am training a GAN. But the generator is a model that learn regression from some prepared input (not random noise). Let’s say we have two model blocks: generator (G) and discriminator (D), and three losses: GAN loss for discriminator (d_loss), GAN loss for generator (g_loss), and regression loss for generator (mse_loss). I would like to update G by combining g_loss and mse_loss.</p>
<p>My question:<br>
It is straightforward to optimize D model, which use GAN loss for discriminator. However, I am not quite sure whether I am doing right. My code is like:</p>
<pre><code class="lang-auto">input = ...
regress_label = ...
G = generator()
D = discriminator()
D_optimizer = torch.optim.Adam(D.parameters())
G_optimizer = torch.optim.Adam(G.parameters())
# train G part
g_out = G(input)
d_g_out = D(g_out)
g_loss = loss(d_g_out, tell_D_real)
mse_loss = mse(g_out, regress_label)
# combine GAN loss and regression
total_g_loss = mse_loss + g_loss     # can I just sum them up ???
total_g_loss.backward()
G_optimizer.step()    # update G only relying on optimizer ???
</code></pre>
<p>Thank you very much !</p>

username: ptrblck
cooked: <p>The code looks correct. Note that <code>lotal_g_loss.backward()</code> would also calculate the gradients for <code>D</code> (if you haven’t set all <code>requires_grad</code> attributes to <code>False</code>), so you would need to call <code>D.zero_grad()</code> before updating it.</p>

username: Max.T
cooked: <p><a class="mention" href="/u/ptrblck">@ptrblck</a> Thank you very much!</p>
<p>I am thinking different architecture of multi-task network should have different BP approach. Here are 2 conditions:</p>
<p>Condition 1:<br>
G(encoder  → decoder) → loss_G<br>
G(encoder → decoder) → D → loss_D<br>
In this condition, we should sum loss_G and loss_D up and do BP to update G.</p>
<p>Condition 2:<br>
G(encoder → decoder) → loss_G<br>
G(encoder) → G(decoder) → loss_G<br>
|_____&gt; D - &gt;loss_D</p>
<p>In this condition, the two tasks like two branches of a tree. Only the G(encoder) is the common part. The two losses should be BP to update G(encoder). But for G(decoder) part, only loss_G should be used.<br>
In this case, can we still feed loss_D + loss_G to optimizer?</p>
<p>Thank you very much!</p>

username: ptrblck
cooked: <p>I’m not quite sure about the second approach, since <code>loss_G</code> is calculated twice.<br>
Would the <code>G(encoder → decoder) → loss_G</code> pipeline be still used or only <code>G(encoder) → G(decoder) → loss_G</code> to calculate <code>loss_G</code>?</p>

username: Max.T
cooked: <p>Sorry for the confusion. In fact, I just want to know how to do BP in multi-task learning in pytorch.</p>
<p>Just forget about the GAN, let’s say the architecture is:</p>
<p>block A - &gt; block B → loss_B<br>
|_____&gt; block C - &gt;loss_C</p>
<p>the model parameters include block A, B and C. optimizer is torch.optim.Adam(model.parameters()), and total loss = loss_B + loss_C. Is this correct to train the multi-task learning model?</p>
<p>Thank you.</p>

username: ptrblck
cooked: <p>Yes, your approach looks correct. Note that:</p>
<pre><code class="lang-python">loss = loss_B + loss_C
loss.backward()
</code></pre>
<p>would accumulate the gradients in <code>block_A</code>, since both forward passes used the parameters in it. I assume that fits your use case and is expected.</p>