Custom loss function for specific label

username: Ed_Moman
cooked: <p>Hello,</p>
<p>I would like to create a custom loss function that takes into consideration only one of the labels.</p>
<p>I have three labels (0, 1, 2) and I would like to consider only 0.</p>
<p>This is my code:</p>
<pre><code class="lang-auto">class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")

        outputs = model(**inputs)
        logits = outputs.get("logits")
        
        label_tensor = labels.view(-1)
        output_tensor = logits.view(-1, self.model.config.num_labels)
        
        mask = label_tensor.eq(0)
        
        loss = torch.mean(torch.abs(torch.masked_select(output_tensor, mask) - torch.masked_select(label_tensor, mask)))
        return (loss, outputs) if return_outputs else loss
</code></pre>
<p>This is not working.</p>
<p>Any ideas?</p>

username: Ed_Moman
cooked: <p>The following works (I do not know whether it makes sense, but at least there are no errors):</p>
<pre><code class="lang-auto">class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")

        outputs = model(**inputs)
        logits = outputs.get("logits")
        
        label_tensor = labels.view(-1)
        output_tensor = torch.argmax(logits, dim=-1)
        
        mask = label_tensor.eq(0)
        
        loss = (torch.masked_select(output_tensor, mask) - torch.masked_select(label_tensor, mask)).float().abs().mean()
        loss.requires_grad = True
        
        return (loss, outputs) if return_outputs else loss
</code></pre>
<p>Thanks a million!</p>

username: ptrblck
cooked: <p>The code looks wrong since it seems you are detaching the <code>logits</code> tensor by calling <code>torch.argmax</code> on it, which will raise a valid error explaining that <code>backward()</code> cannot be called on the loss tensor. It then seems you are trying to fix it by explicitly calling <code>loss.requires_grad = True</code> which will not re-attach the computation graph somehow and only masks the error.<br>
The model parameters will thus not get any valid gradients.</p>

username: Ed_Moman
cooked: <p>Thank you. So perhaps the entire approach does not make sense. Or is there a workaround?</p>
<p>What I am doing right now is using weighted Cross Entropy Loss:</p>
<pre><code class="lang-auto">class_weights = compute_class_weight(class_weight='balanced', classes=np.unique(df['labels']), y=df['labels'])

c_w = np.fromiter((np.ceil(i/min(class_weights)) for i in class_weights), dtype=np.float32)
c_w[np.argmax(c_w)] = np.max(c_w) + 1

class CustomTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        labels = inputs.pop("labels")
        # forward pass
        outputs = model(**inputs)
        logits = outputs.get("logits")
        # compute custom loss (suppose one has 3 labels with different weights)
        loss_fct = nn.CrossEntropyLoss(weight=torch.tensor(c_w, device=model.device))
        loss = loss_fct(logits.view(-1, self.model.config.num_labels), labels.view(-1))
        return (loss, outputs) if return_outputs else loss
</code></pre>
<p>But this is sill fitting primarily the majority class.</p>
<p>However, I am interested in one of the minority classes.</p>
<p>If I increase the weight of the target class above a certain threshold the models just breaks.</p>

username: Ed_Moman
cooked: <p>I am trying upsampling…</p>

username: Ed_Moman
cooked: <p>Upsampling seems to do the trick in this case. The training is much more stable. Then we will see how well the model generalises.</p>
<p>My guess, I do not know if this makes any sense, is that the batch size for evaluation was too small to properly sample the minority classes. But I have very limited VRAM.</p>