How can I use different losses to update different branches respectively, and sum grad together to update master(main) branch

username: RealRain
cooked: <p>I am facing a multi-label problem. My network has 2 branches, which are separated from a conv layers.<br>
There is the sketch of my network</p>
<pre><code class="lang-auto">class MultiLabelDemo(nn.Module):
    def __init__(self):
        super(MultiLabelDemo, self).__init__()
        self.main_block = nn.Sequential(
            nn.Conv2d(in_channels=3, out_channels=96, kernel_size=11, stride=4),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=96, out_channels=96, kernel_size=1, stride=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=96, out_channels=96, kernel_size=1, stride=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)
        )
        self.tail_block1 = nn.Sequential(
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=1, stride=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=1, stride=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)
        )
        self.tail_block2 = nn.Sequential(
            nn.Conv2d(in_channels=96, out_channels=256, kernel_size=5, stride=1, padding=2),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=1, stride=1),
            nn.ReLU(inplace=True),
            nn.Conv2d(in_channels=256, out_channels=256, kernel_size=1, stride=1),
            nn.ReLU(inplace=True),
            nn.MaxPool2d(kernel_size=3, stride=2, ceil_mode=True)
        )

    def forward(self, data):
        x = self.main_block(data)
        y1 = self.tail_block1(x)  
        y2 = self.tail_block2(x)

        return y1, y2

</code></pre>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/9/9283cfe7041b9f507a71cc1415278bccc76e8980.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/9283cfe7041b9f507a71cc1415278bccc76e8980" title="WechatIMG1"><img src="https://discuss.pytorch.org/uploads/default/optimized/2X/9/9283cfe7041b9f507a71cc1415278bccc76e8980_2_281x500.jpeg" alt="WechatIMG1" data-base62-sha1="kU87fa3jmHXmphTlVqElojLMYoM" width="281" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/2X/9/9283cfe7041b9f507a71cc1415278bccc76e8980_2_281x500.jpeg, https://discuss.pytorch.org/uploads/default/optimized/2X/9/9283cfe7041b9f507a71cc1415278bccc76e8980_2_421x750.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/2X/9/9283cfe7041b9f507a71cc1415278bccc76e8980_2_562x1000.jpeg 2x" data-dominant-color="A09D99"><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">WechatIMG1</span><span class="informations">1080×1915 58.3 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>
branch_1 is for 1st label and branch is for 2nd label, so loss_1 is computed  as <code>nn.CrossEntropyLoss(y1, label_1)</code> and loss_2 is computed as <code>nn.CrossEntropyLoss(y2, label_2)</code>.<br>
The whole network has only one optimizer (SGD).(I’m not sure whether it is correct.).</p>
<p><strong>What I want to implement is that</strong> loss_1 is used to update the weights and bias in tail_block1 and loss_2 is used to update the weights and bias in tail_block2. when backprop progress goes to the layer when network was separated, add two grad together and update the remaining main_block. How can I do this, I’ve searched a lot but nothing found. I’d appreciate your kindly help.</p>

username: ptrblck
cooked: <p>Your approach would work out of the box in PyTorch.<br>
You could just add both of your losses and call <code>backward</code> on it.<br>
This will make sure that the tails of your model will use the appropriate gradients, while the common part uses the accumulated gradient.</p>

username: RealRain
cooked: <p>I‘ve seen several similar problems.’<a href="https://discuss.pytorch.org/t/how-to-train-the-network-with-multiple-branches/2152">https://discuss.pytorch.org/t/how-to-train-the-network-with-multiple-branches/2152</a>. So I don’t think it is out of the box of pytorch, and I would like to double check whether it works.</p>

username: ptrblck
cooked: <p>Sure, it’s always good to double check an approach.<br>
In your use case you don’t have to modify anything in a fancy way, but rather just add the different losses together as shown in the thread you’ve linked.</p>
<p>If you are using different criteria for both tails, I would try to make sure the losses have approx. the same range, e.g. by scaling them.</p>

username: RealRain
cooked: <p><img src="https://discuss.pytorch.org/images/emoji/apple/grinning.png?v=5" title=":grinning:" class="emoji" alt=":grinning:"><br>
Thanks for your kindly reply.<br>
I am putting all of my looses into a python list and use a <code>torch.autograd.backward()</code> function to do back propagation. I am running several tests to see if it is correct.</p>

username: ptrblck
cooked: <p>I’m not sure if you will encounter this error, but if you want to call <code>backward</code> sequentially for your losses, you might want to specify <code>retain_graph=True</code> in the <code>.backward()</code> call, as the intermediate variables will be cleared otherwise.</p>

username: justusschock
cooked: <p>If you call <code>backward()</code> on each loss separately you will  for sure encounter this error. The easier way would be to sum up the losses. Both approaches are equivalent since the gradient of multiple backward calls will be accumulated.</p>

username: RealRain
cooked: <p>In fact I am pretty new to pytorch. since I’ve used caffe for several years.  <img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"><br>
I am reading the documentation of pytorch again and again.<img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"><img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"><img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"><img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"><img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"><img src="https://discuss.pytorch.org/images/emoji/apple/joy.png?v=5" title=":joy:" class="emoji" alt=":joy:"></p>

username: RealRain
cooked: <p>Thanks for your suggestion. I’m gonna to check my code.</p>

username: Naruto-Sasuke
cooked: <p>Just add all loss together and call <code>.backward()</code> once! DO NOT respectively call ‘loss1.backward’ and ‘loss2.backward’, since you do not add ‘retain_graph =True’ in the module of ‘main_block’.<br>
A good tutorial of ‘retain_graph=True’ is here <a href="https://github.com/pytorch/tutorials/tree/master/advanced_source" rel="nofollow noopener">https://github.com/pytorch/tutorials/tree/master/advanced_source</a></p>

username: RealRain
cooked: <p>Based on your reply, is the following code right?</p>
<pre><code class="lang-auto">optimizer = optim.SGD(params=my_params_list, lr=....)

loss1_func = nn.CrossEntropyLoss()
loss2_func = nn.CrossEntropyLoss()

loss1 = loss1_func(y1, target1)
loss2 = loss2_func(y2. target2)

loss = sum([loss1, loss2])  # or loss = loss1 + loss2
loss.backward()
optimizer.step()

</code></pre>
<p>from previous reply, I’ve changed my code into:</p>
<pre><code class="lang-auto">optimizer = optim.SGD(params=my_params_list, lr=....)

loss1_func = nn.CrossEntropyLoss()
loss2_func = nn.CrossEntropyLoss()

loss1 = loss1_func(y1, target1)
loss2 = loss2_func(y2. target2)

loss1.backward(retain_gragh=True)
loss2.backward(retain_gragh=True)

optimizer.step()
</code></pre>
<p>Thanks for your suggestion.</p>

username: Naruto-Sasuke
cooked: <p>The first is elegant and correct. ‘loss = loss1+loss2’ is ok.<br>
The second way may consume more memory ‘loss1.backward(retain_graph=True)’ makes the whole branch(including the main_block) stores the intermediate variables</p>

username: RealRain
cooked: <p>Thank you very much for your help. <img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=5" title=":hugs:" class="emoji" alt=":hugs:"><img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=5" title=":hugs:" class="emoji" alt=":hugs:"><img src="https://discuss.pytorch.org/images/emoji/apple/hugs.png?v=5" title=":hugs:" class="emoji" alt=":hugs:"></p>

username: justusschock
cooked: <p>Note that you should zero the gradients before calling backward via <code>optim.zero_grad()</code></p>

username: RealRain
cooked: <p>yeah. I did call <code>optimizer.zero_grad()</code> at the beginning of my training code block.</p>

username: justusschock
cooked: <p>You should do this every time before calling <code>backward</code> somewhere (except you are dealing with RNNs). Otherwise the gradient values of all losses will be accumulated</p>