# How to train multi-branch output network?

**URL:** <https://discuss.pytorch.org/t/how-to-train-multi-branch-output-network/95199>\
**Category:** Uncategorized\
**Created:** [September 4, 2020, 3:16pm UTC](https://discuss.pytorch.org/t/how-to-train-multi-branch-output-network/95199 "2020-09-04T15:16:12Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![MinsooKim](https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/m/c67d28/32.png) [@MinsooKim](https://discuss.pytorch.org/u/MinsooKim)\
**Post date:** [September 4, 2020, 3:16pm UTC](https://discuss.pytorch.org/t/how-to-train-multi-branch-output-network/95199/1 "2020-09-04T15:16:12Z")

</div>

Hello, everyone!

I am working to build multi-branch network.  
Here is what I am working on.

 ![File (4)](https://discuss.pytorch.org/uploads/default/original/3X/8/5/8550c512c607ba20636a78bc52128229d6bdc3b6.jpeg)

Each data sample in the batch has its own class, so each sample goes into a decoder of the same class as itself.  
**How to train this kind of network?**

Here is a training step in my current code.

```auto
params = list(Encoder.parameters()) + list(Decoder1.parameters()) \
                    + list(Decoder2.parameters()) + list(Decoder3.parameters())
optim = torch.optim.Adadelta(params, learning_rate)

encoder_output = Encoder(input)
optim.zero_grad()
loss_1, loss_2, loss_3 = 0, 0, 0
# split batch
if encoder_output[1_batch_ind].shape[0] != 0:
     output1 = Decoder1(encoder_output[1_batch_ind])
     loss_1 = loss_fn(output1, ground_truth[1_batch_ind])
                
if encoder_output[2_batch_ind].shape[0] != 0:
     output2 = Decoder2(encoder_output[2_batch_ind])
     loss_2 = loss_fn(output2, ground_truth[2_batch_ind])
                
if encoder_output[3_batch_ind].shape[0] != 0:
     output3 = Decoder3(encoder_output[3_batch_ind])
     loss_3 = loss_fn(output3, ground_truth[3_batch_ind])
loss = loss_1 + loss_2 + loss_3
loss.backward()
optim.step()             
                

```

What I have expected is each loss\_# is backpropagated through each Decoder#, only and Encoder is trained by a total sum of loss, which means the backpropagation happens like  
loss\_# → Decoder#,  
total\_loss → Encoder.

Does the code above run as I have expected??

---

<div class="post-metadata">

**Author:** ![ptrblck](https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/ptrblck/32/1823_2.png) [@ptrblck](https://discuss.pytorch.org/u/ptrblck)\
**Post date:** [September 7, 2020, 8:27am UTC](https://discuss.pytorch.org/t/how-to-train-multi-branch-output-network/95199/2 "2020-09-07T08:27:06Z")

</div>

Yes, that should work as expected.  
To verify it, you could zero out some losses and check the gradients in each module separately.