Training a multi-task model with a learnable loss component weight

akt42: <p>I’m trying to implement a learnable loss weight for a multi-task setup. (mostly inspired from: <a href="https://discuss.pytorch.org/t/how-to-learn-the-weights-between-two-losses/39681" class="inline-onebox">How to learn the weights between two losses?</a>)</p>
<p>The following snippet is a minimal example of a multi-task model and a mult-task loss with a learnable loss weight.</p>
<p>The issue is the following implementation only updates the loss weight but not the model weights. How can I optimize for both the loss weight and model weights?</p>
<pre><code class="lang-auto">from typing import Dict
import torch
import torch.nn as nn

class Model(nn.Module):

    def __init__(self):
        super().__init__()
        self.classifier = nn.Linear(10, 3)
        self.regressor = nn.Linear(10, 1)
        
    def forward(self, x):
        classifier_op = self.classifier(x)
        regressor_op = self.regressor(x)

        return {
            "classification": classifier_op,
            "regression": regressor_op
        }

class MultiTaskLoss(nn.Module):

    def __init__(self, 
                 loss_mapping: Dict[str, nn.Module] = {"classification": nn.CrossEntropyLoss(), "regression": nn.MSELoss()}):
        super().__init__()
        self.loss_mapping = loss_mapping
        self.loss_weights = nn.Parameter(torch.ones(2))

    def forward(self,
                model_output: Dict[str, torch.Tensor], 
                targets: Dict[str, torch.Tensor]):
        
        losses = []
        
        for task_name, task_op in model_output.items():
            loss_fn = self.loss_mapping[task_name]
            loss_value = loss_fn(task_op, targets[task_name])
            losses.append(loss_value)

        losses = torch.Tensor(losses)

        # loss_weighting
        losses = losses * self.loss_weights

        return losses.sum()

model = Model()
mlt_loss = MultiTaskLoss()

params = list(model.parameters()) + list(mlt_loss.parameters())
optimizer = torch.optim.Adam(params, lr=0.1)

# train
for i in range(3):

    ip = torch.rand((1, 10))

    targets = {
        "classification": torch.Tensor([[1, 0, 0]]),
        "regression": torch.rand([1])
    }

    optimizer.zero_grad()
    model_op = model(ip)
    loss = mlt_loss(model_op, targets)
    loss.backward()
    optimizer.step()
    
    # printing the parameters of the optimizer
    print(f"\nEpoch {i}")
    for t in optimizer.param_groups[0]['params']:
        print(t.shape, t.mean().item())
</code></pre>
<p>The following shows all the learnable parameters applicable to the optimizer:</p>
<pre><code class="lang-auto">for t in optimizer.param_groups[0]['params']:
    print(t.shape, t.requires_grad)

torch.Size([3, 10]) True             &lt; classifer weight
torch.Size([3]) True                 &lt; classifier bias
torch.Size([1, 10]) True             &lt; regressor weight
torch.Size([1]) True                 &lt; regressor bias
torch.Size([2]) True                 &lt; loss weight
</code></pre>
<p>The above snippet prints the mean of the each weight after each epoch (which shows only the loss weight has been updated but not the rest of the parameters - only the last weight is updated in the output after each epoch):</p>
<pre><code class="lang-auto">
Epoch 0
torch.Size([3, 10]) 0.07859472185373306
torch.Size([3]) 0.000278279185295105
torch.Size([1, 10]) -0.04856250435113907
torch.Size([1]) -0.26647862792015076
torch.Size([2]) 0.8999999761581421               

Epoch 1
torch.Size([3, 10]) 0.07859472185373306          &lt; not updated
torch.Size([3]) 0.000278279185295105             &lt; not updated
torch.Size([1, 10]) -0.04856250435113907         &lt; not updated
torch.Size([1]) -0.26647862792015076             &lt; not updated
torch.Size([2]) 0.8023112416267395               &lt; updated

Epoch 2
torch.Size([3, 10]) 0.07859472185373306          &lt; not updated
torch.Size([3]) 0.000278279185295105             &lt; not updated
torch.Size([1, 10]) -0.04856250435113907         &lt; not updated
torch.Size([1]) -0.26647862792015076             &lt; not updated
torch.Size([2]) 0.7098063230514526               &lt; updated

</code></pre>
<p>How can I make both the model and loss weight parameters updated in this setting?</p>

KFrank: <p>Hi akt42!</p>
<aside class="quote no-group quote-modified" data-username="akt42" data-post="1" data-topic="200455" data-full="true">
<div class="title">
<div class="quote-controls"></div>
<img loading="lazy" alt="" width="24" height="24" src="https://discuss.pytorch.org/letter_avatar_proxy/v4/letter/a/258eb7/48.png" class="avatar"> akt42:</div>
<blockquote>
<p>following implementation only updates the loss weight but not the model weights.</p>
<pre><code class="lang-auto">        losses = torch.Tensor(losses)
</code></pre>
</blockquote>
</aside>
<p>This line breaks the computation graph that links <code>losses</code> back to your<br>
model parameters.  So when you backpropagate, the <code>.grad</code> properties<br>
of your model parameters remain <code>None</code> and your optimizer doesn’t update<br>
the parameters.</p>
<p>Change the line to <code>losses = torch.cat (losses)</code> (or something similar,<br>
depending on what shape you need <code>losses</code> to have).</p>
<p>(<code>torch.Tensor()</code> constructs a new tensor that has <code>requires_grad = False</code>.<br>
Furthermore, <code>losses</code> starts out as a list of two single-value tensors (that do<br>
happen to have <code>requires_grad = True</code>).  But they get converted to python<br>
scalars along the way and those have no notion of <code>requires_grad</code>.)</p>
<p>Best.</p>
<p>K. Frank</p>

akt42: <p>Hi Frank,</p>
<p>Thank you! I was able to get it working as follows based on your observation.</p>
<pre><code class="lang-auto">class MultiTaskLoss(nn.Module):

    def __init__(self, 
                 loss_mapping: Dict[str, nn.Module] = {"classification": nn.CrossEntropyLoss(), "regression": nn.MSELoss()}):
        super().__init__()
        self.loss_mapping = loss_mapping
        self.loss_weights = nn.Parameter(torch.ones(2))

    def forward(self,
                model_output: Dict[str, torch.Tensor], 
                targets: Dict[str, torch.Tensor]):
        
        losses = []
        
        for task_name, task_op in model_output.items():
            loss_fn = self.loss_mapping[task_name]
            loss_value = loss_fn(task_op, targets[task_name])
            losses.append(loss_value.unsqueeze(dim=0))

        losses = torch.cat(losses)

        # loss_weighting
        losses = losses * self.loss_weights

        return losses.sum()
</code></pre>