How do I update my network with different losses for each component?

username: Giovanni
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/2/e/2edbb2a285346102bff7863b0fc4376ce86c7632.png" data-download-href="https://discuss.pytorch.org/uploads/default/2edbb2a285346102bff7863b0fc4376ce86c7632" title="螢幕快照 2020-08-24 下午7.42.14"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2edbb2a285346102bff7863b0fc4376ce86c7632_2_690x367.png" alt="螢幕快照 2020-08-24 下午7.42.14" data-base62-sha1="6GwFanozPZQEsv7DfEWSZbHCZEK" width="690" height="367" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2edbb2a285346102bff7863b0fc4376ce86c7632_2_690x367.png, https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2edbb2a285346102bff7863b0fc4376ce86c7632_2_1035x550.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/2/e/2edbb2a285346102bff7863b0fc4376ce86c7632_2_1380x734.png 2x" data-dominant-color="E4F9F9"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">螢幕快照 2020-08-24 下午7.42.14</span><span class="informations">1692×902 37.2 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div><br>
The image is a simplified version of my network. I want to update Network1 with Loss_Network1 and Network2 with Loss_Network2. However, when I do backward(), the grads will be calculated for them regardless (ex: the grads for Loss_Network1 is added onto the .grad of Network2). The only approach I can think of now is:<br>
Loss_Network1.backward()<br>
optimizer_Network2.zero_grad()<br>
copy the .grad for Network1<br>
Loss_Network2.backward()<br>
optimizer_Network1.zero_grad()<br>
copy back the .grad for Network1<br>
optimizer_Network1.step()<br>
optimizer_Network2.step()</p>
<p>and also, doing optimizer_Network1.step() before Loss_Network2.backward() although may seem obvious but isn’t an option because that way it will raise: RuntimeError: one of the variables needed for gradient computation has been modified by an inplace operation</p>
<p>Please let me know if there’s is a simple way to maybe freeze .grad in backward() so that I won’t need to copy the .grad s.    Thanks!!</p>

username: albanD
cooked: <p>Hi,</p>
<p>I’m afraid this is not easy to do. And what you’re doing right now is pretty close to optimal.<br>
The only other approach would be to do two forwards, one for lossnet1 and one for lossnet2 and only update the corresponding net before doing the other’s backward. But this will be slower I’m afraid.</p>

username: Giovanni
cooked: <p>Thanks for the reply! I just thought of another possible solution. Can we do something like stopping the back propagation at some point? like stop it in between Net1 and Net2?</p>

username: Giovanni
cooked: <p>I have just implemented the copy and past grad function.</p>
<pre><code class="lang-auto">def copy_grad(optimizer):
    grad_list = []
    for ppp in optimizer.param_groups[0]['params']:
        gradd = ppp.grad.clone()
        grad_list.append(gradd)

    return grad_list

def paste_grad(optimizer, grad_list):
    for index, ppp in enumerate(optimizer.param_groups[0]['params']):
        ppp.grad = grad_list[index]
        
    return
</code></pre>

username: albanD
cooked: <p>You can, but you will still need 2 calls to backprop. So i don’t think there will be a very large benefit from this.</p>

username: thepurpleowl
cooked: <p>Is there a better/newer way to achieve this, which have been used after last update on this discussion?<br>
<a class="mention" href="/u/alband">@albanD</a></p>