Fine-Tune BERT with two Classification Heads "next to each other"?
I am currently working on a project to fine-tune BERT models on a multi-class classification task with the goal to classify job ads into some broader categories like “doctors” or “sales” viaAutoModelForSequenceClassification(which works quite well). Now I am wondering, whether it would be possible to add a second classification head “next” to the first one (not in sequence) to classify the minimum educational level that is required for the job. I imagine that each head is directly connected to the pooler output an then makes a prediction independent of the other’s prediction. I think my use-case is slightly different than a multi-label classification since both labels describe different aspects of the job ad.Similar to this CV example:What is a multi-headed model? And what exactly is a ‘head’ in a model?I hope it’s not total nonsense that I’m asking hereGreetings,David
2021-09-14T15:19:06Z
Sure, you can just use any default model, e.g.BertModeland add your custom classification heads on top of that. Have a look at the existing iclassification implementation. You can basically duplicate that, but add another classifier layer. Of course you’ll also have to adapt the forward method accordingly.github.comhuggingface/transformers/blob/c1e47bf4fe1d9de06d774cc2c24ec5a93461c5a5/src/transformers/models/bert/modeling_bert.py#L1481)@add_start_docstrings("""Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of the pooledoutput) e.g. for GLUE tasks.""",BERT_START_DOCSTRING,)class BertForSequenceClassification(BertPreTrainedModel):def __init__(self, config):super().__init__(config)self.num_labels = config.num_labelsself.config = configself.bert = BertModel(config)classifier_dropout = (config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob)self.dropout = nn.Dropout(classifier_dropout)
2021-09-17T08:35:44Z
@BramVanroythanks for the swift reply! So I tried to get my head around the first part, which is setting up the custom model but I am not sure if I fully understand what I have to do (still working in my pytorch and transformer skills…). These are my current adjustments to the class:class CustomBertForSequenceClassification(BertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config
      
        model_name = 'bert-base-german-cased'
        self.bert = BertPreTrainedModel.from_pretrained(model_name) #BertModel(config)

        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )

        self.dropout = nn.Dropout(classifier_dropout)

        self.classifier_job_type = nn.Linear(config.hidden_size, config.num_labels)
        self.classifier_edu_level = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()I added a linear layersself.classifier_job_typeandself.classifier_edu_level. The model looks like this now:CustomBertForSequenceClassification(
  (bert): BertPreTrainedModel()
  (dropout): Dropout(p=0.1, inplace=False)
  (classifier_job_type): Linear(in_features=768, out_features=2, bias=True)
  (classifier_edu_level): Linear(in_features=768, out_features=2, bias=True)
)So maybe you can tell me whether I am on the right track…Something I don’t quite get is why the out_features are 2 are they coming from the config? How can I adjust the number of labels for both classification layers?Thanks in advance!David
2021-09-17T08:45:48Z
Yes, you are on the right track! I encourage you to go look at the PyTorch documentation to better understand what you are doing. You’ll find that the second argument to nn.Linear is the output dimension.out_features. You have implemented this in such a way that it gets the default values from config. Instead, you can change those output features as you wish. E.g.self.classifier_job_type = nn.Linear(config.hidden_size, 5)
self.classifier_edu_level = nn.Linear(config.hidden_size, 3)It may seem from that printed model that those classification layers are sequential in the model, but do not worry, that is not the case. It simply represents the order in which you defined the layers. The “actual” data flow order is controlled by theforwardmethod, that you still have to implement. It should go like this (in wors, so that you can implement it yourself):Push the encoded inputs through self.bertPush the output of self.bert through a dropout and then through the first classification layerPush the output of self.bert through a dropout and then through the second classification layerCalculate loss of output first classification and of output in second classification layer and then sum the lossDuring training you then can do a single backward pass