Weighted Loss in BertForTokenClassification
Author: krishanudb
Pull request: no
Labels: name: Feature request

# 🚀 Feature request

BertForTokenClassification models can compute cross entropy loss currently is only weighted. The option to have different weights for different classes can be useful in several use cases, including but not restricted to the problem of unbalanced output classes

## Motivation

Right now, although BertForTokenClassification models can compute cross entropy loss during the forward pass, there is no explicit way of weighting the different classes, which seems like a useful feature, as sequence tagging tasks often have unbalanced classes. I ran into the above problem during solving a academic problem. I looked at the code for the BertForTokenClassification model, and found that it should be quite easy to implement

## Your contribution

Not sure if I can help, coz I am not really familiar with the codebase. However, I can point out how and where to add code to implement the weighted loss very easily

Comments:
In PyTorch, [`nn.CrossEntropyLoss`](https://pytorch.org/docs/stable/generated/torch.nn.CrossEntropyLoss.html) has an optional `weight` parameter which you can specify. This should be a 1D Tensor assigning a weight to each of the classes.

So if you want `BertForTokenClassification` with a weighted cross entropy loss, you can simply replace [this line](https://github.com/huggingface/transformers/blob/c60e0e1ee45f4bf1017736b146c51729f120bb83/src/transformers/models/bert/modeling_bert.py#L1685) by a weighted loss. For example, you can define it as follows (I just copied the relevant code from `modeling_bert.py` and slightly adapted the cross entropy loss):

```
class BertForTokenClassification(BertPreTrainedModel):

    _keys_to_ignore_on_load_unexpected = [r"pooler"]

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

        self.bert = BertModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()

    @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="bert-base-uncased",
        output_type=TokenClassifierOutput,
        config_class=_CONFIG_FOR_DOC,
    )
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
        return_dict=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
            Labels for computing the token classification loss. Indices should be in ``[0, ..., config.num_labels -
            1]``.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.bert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )

        sequence_output = outputs[0]

        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)

        loss = None
        if labels is not None:
            weights = torch.tensor([0.6, 0.3, 0.1])
            loss_fct = CrossEntropyLoss(weights=weights)
            # Only keep active parts of the loss
            if attention_mask is not None:
                active_loss = attention_mask.view(-1) == 1
                active_logits = logits.view(-1, self.num_labels)
                active_labels = torch.where(
                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
                )
                loss = loss_fct(active_logits, active_labels)
            else:
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return TokenClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )
```


@NielsRogge 
You are right. I had done exactly this in my local (huggingface) transformers codebase.
Worked as expected.

I think this would be a useful feature if huggingface models come with it.

Oh ok so you want this to be an added feature. Not sure if this is possible. @LysandreJik what do you think?

Hi, thanks for opening an issue! The losses in the models are not made to be completely customizable, but to be the most common loss used in most cases; we favor simplicity here.

This is because defining your custom loss in a PyTorch model is very simple: when you do not pass the labels to your model, then you retrieve the model logits. You can then define a loss (and customize it as you wish!) and compute its value using these logits and your labels.

However, this is not the first time this feature has been requested, and we could probably come up with an implementation that wouldn't complexify the code-base too much. If we see more of this request we'll take a deeper look at how to implement it.

Here's a past issue discussing the same/similar: https://github.com/huggingface/transformers/issues/7024

cc @sgugger @patrickvonplaten

I'm not sure whether it's a good idea to add such functionality to `modeling_bert.py` - there are too many possibilities. I think it could very well be added to the examples though.

Yes. Unfortunately, there are too many of these possibilities. 
Most users who are familiar with PyTorch can anyway make necessary changes to their local codebase quite easily.
Thanks.

Also see [this example in the documentation](https://huggingface.co/transformers/main_classes/trainer.html) (scroll a tiny bit down to the first example showing a subclass of `Trainer`) on how to change just the loss computation while using a model with `Trainer`.

Hi everyone,

I am a student and therefore not yet very familiar with the way issues report work on git, so I aplogize in advance if this is not the proper place to post this message.
I've stumbled onto an error when using the aforementioned method for designing a custom loss function.
My code is the following
```
config = AutoConfig.from_pretrained("bert-base-cased", num_labels=2, finetuning_task="SST-2")

# Test with modified trainer for weighted CrossEntropyLoss

model = AutoModelForSequenceClassification.from_pretrained(
        "dmis-lab/biobert-base-cased-v1.1",
        from_tf=False,
        config=config)

from torch import FloatTensor

classDistribution_raw = [97, 3]
classDistribution = [0.8, 0.2]
normedWeights = [1 - (x / sum(classDistribution)) for x in classDistribution]
normedWeights = FloatTensor(normedWeights).cuda()

from torch.nn import CrossEntropyLoss

class MyTrainer(Trainer):
    def compute_loss(self, model, inputs, return_outputs=False):
        
        if "labels" in inputs:
          labels = inputs.pop("labels")
        
        outputs = model(**inputs)
        logits = outputs.logits
        loss_function = CrossEntropyLoss(weight = normedWeights)

        if self.args.past_index >= 0:
            self._past = outputs[self.args.past_index]

        if labels is not None:
            loss = loss_function(logits, labels)
        else:
            # We don't use .loss here since the model may return tuples instead of ModelOutput.
            loss = outputs["loss"] if isinstance(outputs, dict) else outputs[0]

        return (loss, outputs) if return_outputs else loss

trainer = MyTrainer(
        model=model,
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
        compute_metrics=compute_metrics_fn,
        tokenizer=tokenizer,
    )

```
And when I try to train the model using trainer.train(), i get the following error
'NoneType' object has no attribute 'detach'

There is probably something wrong with the way I customized the loss function but I can't find where.

Best regards,
Arthur


> Hi, thanks for opening an issue! The losses in the models are not made to be completely customizable, but to be the most common loss used in most cases; we favor simplicity here.
> 
> This is because defining your custom loss in a PyTorch model is very simple: when you do not pass the labels to your model, then you retrieve the model logits. You can then define a loss (and customize it as you wish!) and compute its value using these logits and your labels.
> 
> However, this is not the first time this feature has been requested, and we could probably come up with an implementation that wouldn't complexify the code-base too much. If we see more of this request we'll take a deeper look at how to implement it.
> 
> Here's a past issue discussing the same/similar: #7024
> 
> cc @sgugger @patrickvonplaten

@sgugger @LysandreJik @NielsRogge 

I want to put a +1 on this feature request. Datasets with imbalanced datasets would benefit a lot from custom loss functions. And this shouldn't be a complex add (should just be one more kwarg?).