How to train a network with different branches (not parallel)?

username: whisht
cooked: <p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/2X/2/27ee6300ea72fbff203b97890c1e3949a5aa12fc.png" data-download-href="https://discuss.pytorch.org/uploads/default/27ee6300ea72fbff203b97890c1e3949a5aa12fc" title="image"><img src="https://discuss.pytorch.org/uploads/default/original/2X/2/27ee6300ea72fbff203b97890c1e3949a5aa12fc.png" alt="image" data-base62-sha1="5HfmtHPLFCHtnQOhjzhpzKcgjwM" width="690" height="268" data-dominant-color="FAFAFA"><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">image</span><span class="informations">712×277 3.82 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>
As you can see. The middle outputs of one branche is the input of another branches.</p>
<pre data-code-wrap="python"><code class="lang-python">class MyNet(nn.Module):
   def __init__(self):
        self.convnet = CNN()
        self.encoder1 = Encoder()
        self.discriminator = Discriminator()
        self.encoder2 = Encoder()
        self.classifier = Classifier()
   def forward(self,x):
       #...
       x = self.convnet(x)
       out = self.encoder1(x)
       out1 = self.discriminator(out)
       out2 = self.encoder(2)
       out2 = torch.cat(out,out2)
       out2 = self.classifier(out2)
       return out1,out2
</code></pre>
<p>May be I should write the Net separately:</p>
<pre data-code-wrap="python"><code class="lang-python">class Branch1(nn.Module):
   def __init__(self):
       # ...
        self.convnet = CNN()
        self.encoder = Encoder()
        self.discriminator = Discriminator()
   def forward(self,x):
       x = self.convnet(x)
       out = self.encoder(x)
       out1 = self.discriminator(out)
       return out,out1
</code></pre>
<pre data-code-wrap="python"><code class="lang-python">class Branch2(nn.Module):
   def __init__(self):
       # ...
        self.convnet = CNN()
        self.encoder = Encoder()
        self.classifier = classifier()
   def forward(self,x,branch):
       x = self.convnet(x)
       x= torch.cat(x,branch)
       out = self.encoder(x)
       out = self.classifier(out)
       return out,out1
</code></pre>
<p>Then, How to handle the 2 different  parts with different loss?</p>

username: el_samou_samou
cooked: <p>Hi,</p>
<p>I think the first option is the right one (i corrected a typo):</p>
<pre><code class="lang-auto">class MyNet(nn.Module):
   def __init__(self):
        self.convnet = CNN()
        self.encoder1 = Encoder()
        self.discriminator = Discriminator()
        self.encoder2 = Encoder()
        self.classifier = Classifier()
   def forward(self,x):
       #...
       x = self.convnet(x)
       out = self.encoder1(x)
       out1 = self.discriminator(out)
       out2 = self.encoder2(x)
       out2 = torch.cat(out,out2)
       out2 = self.classifier(out2)
       return out1,out2
</code></pre>
<p>and regarding loss you just write your loss with respect to both output like:</p>
<pre><code class="lang-auto">loss = criterion1(out1, ground_truth1) + alpha * criterion2(out2, ground_truth2)
loss.backward()
opt.step()
</code></pre>
<p>And it will work fine!</p>
<p>Hope this help.</p>