Correctly Backpropagating Losses to different Models

username: Symo
cooked: <p>I am currently working on implementing a Vision Transformer Architecture, which consists of a Shared Encoder and three Decoders for three different tasks, primarily focusing on white balancing. Each of these four components has been implemented as its own <code>nn.Module</code>.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/5/954e53b6e1802f00840048c153b21fb3eb452f9c.jpeg" data-download-href="https://discuss.pytorch.org/uploads/default/954e53b6e1802f00840048c153b21fb3eb452f9c" title="TransCC Architecture"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/954e53b6e1802f00840048c153b21fb3eb452f9c_2_690x295.jpeg" alt="TransCC Architecture" data-base62-sha1="liOXok0CUw1iAQbmAqeecjStb5y" width="690" height="295" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/954e53b6e1802f00840048c153b21fb3eb452f9c_2_690x295.jpeg, https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/954e53b6e1802f00840048c153b21fb3eb452f9c_2_1035x442.jpeg 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/5/954e53b6e1802f00840048c153b21fb3eb452f9c_2_1380x590.jpeg 2x" data-dominant-color="CFCFCC"><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">TransCC Architecture</span><span class="informations">1544×661 115 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>During the training phase, all four components are trained together. The Primary Decoder produces three separate losses (L1, L_ang, L_surf), while the two Auxiliary Decoders each have one loss (L_A and L_E). The loss of the Encoder should be the sum of these five losses (+ L_cont that I’ve not implemented yet).</p>
<p>My challenge lies in correctly backpropagating these losses and informing PyTorch which model each loss should be backpropagated to. Attempting to use multiple <code>.backward()</code> calls results in an error. I have come across suggestions to simply sum all the losses together, allowing the method to automatically handle the backpropagation. However, I do not fully understand this approach and I am unsure if it is effectively assigning the losses as intended.</p>
<p>Here is the code I am currently using for the backward pass within the Training function:</p>
<pre><code class="lang-auto">    with torch.autocast(device.type if device.type != 'mps' else 'cpu', enabled=amp):
                         pred, achromatic_map, pred_hed = get_decoders_outputs(modules, *modules["encoder"](input))
                         total_loss = get_total_loss(input, pred, gt, achromatic_map, pred_hed, gt_hed, loss_weights)

    # Loss
    set_zero_grad(optimizers)
    grad_scaler.scale(total_loss).backward()
    update_modules(modules, optimizers, grad_scaler)
</code></pre>
<p>Additionally, here is the code for computing the losses:</p>
<pre><code class="lang-auto">def get_primary_loss(input, pred, gt, loss_weights):
    loss_score = 0 

    loss_score += L1_loss(pred, gt, loss_weights[0])

    loss_score += angular_loss(input, pred, gt, loss_weights[1])

    loss_score += surf_loss(pred, gt, loss_weights[2])

    return loss_score

def get_total_loss(input, pred, gt, achromatic_map, pred_hed, gt_hed, loss_weights):

    primary_loss = get_primary_loss(input, pred, gt, loss_weights[2:5])
    achromatic_loss = achromatic_pixel_detection_loss(gt, achromatic_map, loss_weights[0])
    edge_loss = edge_detection_loss(pred_hed, gt_hed, loss_weights[1])
    total_loss = primary_loss + edge_loss + achromatic_loss

    return total_loss
</code></pre>

username: ptrblck
cooked: <p>Summing the losses and calling <code>backward</code> on it will backpropagate through each of the summands and should work. In the end the addition of the losses is a differentiable operation which Autograd understands.<br>
If you call <code>backward()</code> multiple times I assume you are running into an error claiming the computation graph was already freed which would indicate that the losses share some common layers/parameters.<br>
In this case, the first <code>backward</code> call will free the intermediate activations needed to compute the gradients. You could work around this issue by using <code>retain_graph=True</code> in the first <code>backward</code> calls and set it back to <code>False</code> in the last one, but summing the losses should also work.</p>