# Fine-tuning BERT with multiple classification heads
**URL:**
**Author:**  [@SaraAmd](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FSaraAmd)%5C
**Post date:** [October 14, 2022, 4:27pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F1 "2022-10-14T16:27:51Z")
I need to train a model that has the same backbone such as BERT as a feature extractor and use multiple classification heads. This scenario is similar to multi-task learning but all the tasks are classification tasks so I will need multiple classification heads. does anyone have any similar notebook code that I can start with?
---
**Author:**  [@shensmobile](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fshensmobile)%5C
**Post date:** [October 24, 2022, 6:35pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F2 "2022-10-24T18:35:27Z")
HI SaraAmd, Iām looking at doing something very similar. I currently have ~15 classification models that all use the same language model. I have a feeling that I will have to write my own forward() function in the end. Iām curious if you made any progress, and if so, I was wondering if you could provide me some wisdom on the subject. Iām wondering if I need to re-train all of these models at the same time, or if I can extract the classification heads from my existing models and pack them together into a new model with a custom forward() function. If you found any existing literature that could put me on the right path, thatād be awesome too š
---
**Author:**  [@SaraAmd](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2FSaraAmd)%5C
**Post date:** [October 24, 2022, 7:01pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F3 "2022-10-24T19:01:26Z")
Hi, Unfortunately, I havenāt made any progress and I havenāt found any literature. But I also assume that the forward function needs to be implemented as well as the loss function for each classification head (all will have cross entropy but how I am not sure). We can brainstorm together to move forward. My goal is to have one BERT model as the feature extractor and then add n-number of classification heads on top of it to train the model is this your goal as well? I feel like you want to train different models separately.
---
**Author:**  [@shensmobile](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fshensmobile)%5C
**Post date:** [October 24, 2022, 7:49pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F4 "2022-10-24T19:49:34Z")
Yes that is exactly my goal. Currently I am just barely able to fit most of my models onto my GPU at inference time, but the goal will be to have a smaller footprint on my GPU so that I can process data in batches instead. I donāt mind writing the forward function, I will likely begin tinkering with it next week. I just need to figure out how to extract the classification heads and pickle them so that I can import them into one large model.
As I mentioned, I have all my models trained separately already š But if someone has pre-written code to do multi-head training all at once, I donāt mind re-training if it means re-using code. Iām happy either way. Iāll share code with you as soon as I start putting pen to paper.
---
**Author:**  [@swtb](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fswtb)%5C
**Post date:** [January 10, 2024, 4:15pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F5 "2024-01-10T16:15:27Z")
Did you folks managed to find a solution to this in the end? Iāve tried overwriting several classes to handle multihead classification for coarse and fine grained labels.
---
**Author:**  [@shensmobile](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fshensmobile)%5C
**Post date:** [January 16, 2024, 5:34pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F6 "2024-01-16T17:34:13Z")
No, I unfortunately did not get around to tackling this. I ended up brute forcing it with a bigger GPU with more memory. I still want to solve this in the long term though, but I think the only answer is to write a custom forward() function that creates multiple parallel output linear layers that each generate one set of prediction/classifications.
---
**Author:**  [@swtb](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fswtb)%5C
**Post date:** [January 17, 2024, 9:01am UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F7 "2024-01-17T09:01:24Z")
I think you are right. In fact, I am able to get the model training by overwriting the following:
- compute\_loss
- forward
- the model itself, to add the layers
- datacollator
- get\_train\_dataloader/ eval\_dataloder
I get it to train but I cannot get it to compute any metrics, for some reason the HF Trainer is skipping over compute\_metrics at eval time.
---
**Author:**  [@nielsr](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fnielsr)%5C
**Post date:** [January 17, 2024, 6:17pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F8 "2024-01-17T18:17:57Z")
Have you considered Google? Itās a pretty good search engine š
> **[Multi-Task Learning with Transformers: Transformers with Multiple Prediction...](https://medium.com/@shahrukhx01/multi-task-learning-with-transformers-part-1-multi-prediction-heads-b7001cf014bf)**
>
> Since the inception of transformer architecture and the birth of hugging face transformers package of Natural Language Processing as aā¦
---
**Author:**  [@swtb](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fswtb)%5C
**Post date:** [January 19, 2024, 12:17pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F9 "2024-01-19T12:17:32Z")
I do appreciate the article, it may be useful for the other posters but it doesnāt give me any clues on getting the compute metrics function to work during the evaluation step. So far I have narrowed it down to the all\_labels variable being set to None prior to this function being called in the eval\_loop.
On the whole, as a senior community member, I think its worth reflecting on how a response like this creates a culture in which newer community members lose the confidence to even ask questions.
Thank you for sharing the article.
---
**Author:**  [@nielsr](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fnielsr)%5C
**Post date:** [January 19, 2024, 12:40pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F10 "2024-01-19T12:40:14Z")
Hi,
Sorry for the earlier reply! One should be able to ask any question.
One can add a `compute_metrics` function to the blog post above. The `compute_metrics` function takes in a named tuple, which one can split into labels and logits, e.g.:
```auto
def compute_metrics(eval_preds):
metric = evaluate.load("glue", "mrpc")
logits, labels = eval_preds
predictions = np.argmax(logits, axis=-1)
return metric.compute(predictions=predictions, references=labels)
```
Hence, in case of a multi-task model as shown in the blog above, the logits depend on the task head being used. To circumvent this, one could add an attribute to the modelās init called ātask\_nameā, which you can change depending on which task you want to evaluate for. Then you can call `trainer.evaluate()` depending on which task is set.
---
**Author:**  [@swtb](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Fu%2Fswtb)%5C
**Post date:** [January 19, 2024, 1:22pm UTC](/static-proxy?url=https%3A%2F%2Fdiscuss.huggingface.co%2Ft%2Ffine-tuning-bert-with-multiple-classification-heads%2F24453%2F11 "2024-01-19T13:22:13Z")
Thanks again, this is helpful for me. Am I right in thinking that the task\_name variable effectively chooses the right head task dependently? If so, what if a single task requires two heads? Such as in my own case I need to train the model to give a coarse and a fine-grained output.
For some context, here is my loss function:
```auto
def compute_loss(self, model, inputs, return_outputs=False):
labels = inputs.pop('labels')
labels_coarse = labels[0][0]
labels_fine = labels[0][1]
outputs = model(**inputs)
logits_coarse, logits_fine = outputs
loss_fct = torch.nn.CrossEntropyLoss()
loss_coarse = loss_fct(logits_coarse.view(-1, num_labels_coarse), labels_coarse.view(-1))
loss_fine = loss_fct(logits_fine.view(-1, num_labels_fine), labels_fine.view(-1))
loss = (loss_coarse + loss_fine) / 2
return (loss, outputs) if return_outputs else loss
```