Different loss functions update different parts of network

Yongjie_Shi: <p>I have two networks, for example, netA and netB. However, I want to use different loss functions to update different networks. For example, I want to use loss1 to update netA, and use loss2 to update netB.<br>
I have defined different optimizers for netA and netB. Each optimizer optimizes only the parameters of the corresponding network.<br>
The total pipeline is shown in the following diagram.<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/6/3/63b99611a4c33702f2f7255585b05fb55b61511e.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/63b99611a4c33702f2f7255585b05fb55b61511e" title="微信图片编辑_20210728145659"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63b99611a4c33702f2f7255585b05fb55b61511e_2_690x146.jpeg" alt="微信图片编辑_20210728145659" data-base62-sha1="eecSWvy0qlC4j38PCzy9bqRwceW" width="690" height="146" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63b99611a4c33702f2f7255585b05fb55b61511e_2_690x146.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63b99611a4c33702f2f7255585b05fb55b61511e_2_1035x219.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/6/3/63b99611a4c33702f2f7255585b05fb55b61511e_2_1380x292.jpeg 2x" data-dominant-color="858480"><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">微信图片编辑_20210728145659</span><span class="informations">2766×588 48.8 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 optimized pipeline I came up with is shown below：<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/5/c/5c7fd216297e81f1bfb41d235d27503797ebecc6.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/5c7fd216297e81f1bfb41d235d27503797ebecc6" title="微信图片编辑_20210728145727"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c7fd216297e81f1bfb41d235d27503797ebecc6_2_361x500.jpeg" alt="微信图片编辑_20210728145727" data-base62-sha1="dchMt5RlkYxdqMTs4FcYFZNfHRc" width="361" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c7fd216297e81f1bfb41d235d27503797ebecc6_2_361x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c7fd216297e81f1bfb41d235d27503797ebecc6_2_541x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/5/c/5c7fd216297e81f1bfb41d235d27503797ebecc6_2_722x1000.jpeg 2x" data-dominant-color="9F9D99"><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">微信图片编辑_20210728145727</span><span class="informations">1855×2566 144 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>
What I want to ask is, is the above process correct?</p>

Yongjie_Shi: <p>My concern is that, is it necessary to forward the data twice, the first time to optimize netA, the second time to optimize netB?</p>

KFrank: <p>Hi Yongjie!</p>
<aside class="quote no-group" data-username="Yongjie_Shi" data-post="1" data-topic="127910" 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/y/e480ec/48.png" class="avatar"> Yongjie_Shi:</div>
<blockquote>
<p>What I want to ask is, is the above process correct?</p>
</blockquote>
</aside>
<p>Your process it mostly correct.  You need to use <code>retain_graph = True</code><br>
in your first <code>backward()</code> call:</p>
<pre data-code-wrap="python"><code class="lang-python">loss1.backward (retain_graph = True)
</code></pre>
<aside class="quote no-group">
<blockquote>
<p>is it necessary to forward the data twice</p>
</blockquote>
</aside>
<p>No, one forward pass is enough.  (You do, of course, have to calculate<br>
both <code>loss1</code> and <code>loss2</code>.)  But you do need two <code>backward()</code> passes.</p>
<p>(If you <em>weren’t</em> to use <code>retain_graph = True</code> in your first <code>backward()</code><br>
pass, you would have to perform a second forward pass to rebuild the<br>
computation graph for use by the second <code>backward()</code> pass.)</p>
<p>Here is an illustrative script:</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
print (torch.__version__)

weights1A = None   # net1 weights after net1 half-update
weights2A = None   # net2 weights after net2 half-update
weights1B = None   # net1 weights after full update
weights2B = None   # net2 weights after full update

for  updateMethod in range (3):   # update model with two "half-updates" vs. one full update
    torch.manual_seed (2021)   # make initialization reproducible

    net1 = torch.nn.Linear (2, 2)   # create new, freshly-initialized model
    net2 = torch.nn.Linear (2, 2)
    input = torch.randn (2)
    target1 = torch.randn (2)
    target2 = torch.randn (2)
    opt1 = torch.optim.SGD (net1.parameters(), lr = 0.1)
    opt2 = torch.optim.SGD (net1.parameters(), lr = 0.2)

    if  updateMethod == 0:   # net1 half-update
        output = net2 (torch.tanh (net1 (input)))
        loss1 = torch.nn.MSELoss() (output, target1)
        loss1.backward()
        opt1.step()
        weights1A = net1.weight.detach().clone()

    elif  updateMethod == 1:   # net2 half-update
        output = net2 (torch.tanh (net1 (input)))
        loss2 = torch.nn.MSELoss() (output, target2)
        loss2.backward()
        opt2.step()
        weights2A = net2.weight.detach().clone()

    elif  updateMethod == 2:   # perform full update with a single forward pass
        output = net2 (torch.tanh (net1 (input)))
        loss1 = torch.nn.MSELoss() (output, target1)   # net1 update
        loss1.backward (retain_graph = True)   # keep graph for net2 loss2.backward()
        opt1.step()
        weights1B = net1.weight.detach().clone()
        opt2.zero_grad()   # zero out net2 grads populated by loss1.backward()
        loss2 = torch.nn.MSELoss() (output, target2)   # net2 update
        loss2.backward()
        opt2.step()
        weights2B = net2.weight.detach().clone()

print ('weights1A = ...')
print (weights1A)
print ('weights1B = ...')
print (weights1B)
print ('check weights1:', torch.allclose (weights1A, weights1B))
1
print ('weights2A = ...')
print (weights2A)
print ('weights2B = ...')
print (weights2B)
print ('check weights2:', torch.allclose (weights2A, weights2B))
</code></pre>
<p>Here is its output:</p>
<pre><code class="lang-plaintext">1.9.0
weights1A = ...
tensor([[-0.4652,  0.0495],
        [ 0.3326,  0.2997]])
weights1B = ...
tensor([[-0.4652,  0.0495],
        [ 0.3326,  0.2997]])
check weights1: True
weights2A = ...
tensor([[-0.6444,  0.6545],
        [-0.2909, -0.5669]])
weights2B = ...
tensor([[-0.6444,  0.6545],
        [-0.2909, -0.5669]])
check weights2: True
</code></pre>
<p>Best.</p>
<p>K. Frank</p>