Model with two output branches optimization

username: sbht_sbht
cooked: <p>How can i adjust the weights in a model with two outputs like this?<br>
<div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/d/0/d04e2623d3bdebbdc5333fa07fbb9698f7c89aac.png" data-download-href="https://discuss.pytorch.org/uploads/default/d04e2623d3bdebbdc5333fa07fbb9698f7c89aac" title="изображение_2022-10-30_131612026"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d04e2623d3bdebbdc5333fa07fbb9698f7c89aac_2_597x500.png" alt="изображение_2022-10-30_131612026" data-base62-sha1="tIKKkxHz9W9QOCW5K5sF4cC3HQw" width="597" height="500" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/d/0/d04e2623d3bdebbdc5333fa07fbb9698f7c89aac_2_597x500.png, https://discuss.pytorch.org/uploads/default/original/3X/d/0/d04e2623d3bdebbdc5333fa07fbb9698f7c89aac.png 1.5x, https://discuss.pytorch.org/uploads/default/original/3X/d/0/d04e2623d3bdebbdc5333fa07fbb9698f7c89aac.png 2x" data-dominant-color="F3F3F3"><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">изображение_2022-10-30_131612026</span><span class="informations">757×633 34.1 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></p>
<p>This is how I implemented it:</p>
<pre><code class="lang-auto">class UNet(nn.Module):
    def __init__(self):
        super(UNet, self).__init__()
        self.down = ConvNetContractingPath()
        self.up = ConvNetExpansivePath()

        self.shuffle_layer = nn.Sequential(
            nn.Conv2d(in_channels=16, out_channels=16, kernel_size=1),
            nn.BatchNorm2d(16)
        )
        self.rotation_layer = nn.Sequential(
            nn.MaxPool2d(2, stride=2),
            nn.Conv2d(in_channels=16, out_channels=4, kernel_size=28, stride=28),
            nn.BatchNorm2d(4)
        )

    def forward(self, x):
        conv5, conv4, conv3, conv2, conv1 = self.down(x)
        up_conv9 = self.up(conv5, conv4, conv3, conv2, conv1)
        shuffle_output = self.shuffle_layer(up_conv9)
        rotation_output = self.rotation_layer(up_conv9)

        return shuffle_output, rotation_output
</code></pre>
<p>I tried to use this method: <a href="https://discuss.pytorch.org/t/a-model-with-multiple-outputs/10440/2" class="inline-onebox">A model with multiple outputs - #2 by ptrblck</a><br>
Calculate the error of each prediction, then sum them up and optimize the model by the sum of the losses. But after 100 epochs the loss value remained almost unchanged.</p>
<pre><code class="lang-auto">criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters(), lr=0.0001)
scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=15, threshold=0.00001)

for epoch in range(num_epochs):
    print(f'Epoch: {epoch+1}/{num_epochs}:')
    for i, (images, shfl_labels, rot_labels) in enumerate(train_loader):

        # Origin shape: image [batch size, 3, 224, 224]
        #               shfl_labels [batch size, 224, 224]
        #               rot_labels [batch size, 4, 4]


        # Load data in GPU
        images = images.to(device)
        shfl_labels = shfl_labels.to(device)
        rot_labels = rot_labels.to(device)

        # Forward pass
        shfl_output, rot_output = model(images)
        loss_shfl = criterion(shfl_output, shfl_labels.long())
        loss_rot = criterion(rot_output, rot_labels.long())
        loss = loss_shfl + loss_rot

        #Backward pass and optimize
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        if (i+1) % 25 == 0:
            print(f'\t Step: {i+1}/{n_total_steps}, Train_loss: {loss.item():.4f}')
</code></pre>
<p>I don’t understand how to optimize the weights of conv2d_15 and conv2d_14 by the amount of losses , since conv2d_15 and conv2d_14 only affect one output?<br>
Maby I should use two optimizers instead of one?</p>

username: srishti-git1110
cooked: <aside class="quote no-group" data-username="sbht_sbht" data-post="1" data-topic="164732">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sbht_sbht/48/46999_2.png" class="avatar"> sbht_sbht:</div>
<blockquote>
<p>how to optimize the weights of conv2d_15 and conv2d_14 by the amount of losses , since conv2d_15 and conv2d_14 only affect one output?</p>
</blockquote>
</aside>
<p>That shouldn’t be a problem since your final <code>loss</code> is the sum of losses from both the outputs, which means that the weights of both conv2d_15 and conv2d_14 are in the graph of <code>loss</code>.</p>
<aside class="quote no-group" data-username="sbht_sbht" data-post="1" data-topic="164732">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/user_avatar/discuss.pytorch.org/sbht_sbht/48/46999_2.png" class="avatar"> sbht_sbht:</div>
<blockquote>
<p>But after 100 epochs the loss value remained almost unchanged.</p>
</blockquote>
</aside>
<p>Maybe reduce the argument <code>patience</code> and see if an automatic change of <code>lr</code> by the scheduler helps. You could also tweak the threshold to only focus on significant changes.</p>

username: sbht_sbht
cooked: <p>You were right<br>
Training did not occur due to an incorrectly configured sheduler<br>
Thanks for your help!</p>