Mismatched size on BertForSequenceClassification from Transformers and multiclass problem

Viewed 1778

I just trained a BERT model on a Dataset composed by products and labels (departments) for an e-commerce website. It's a multiclass problem. I used BertForSequenceClassification to predict the department for each product. I split it in train and evaluation, I used dataloader from pytorch, and I've got a good score with no overfit.

Now I want to try it on a new Dataset to check how it works on unseen data. But I can't achieve to load the model and apply on the new Dataset. I get the following error:

RuntimeError: Error(s) in loading state_dict for BertForSequenceClassification:
    size mismatch for classifier.weight: copying a param with shape torch.Size([59, 1024]) from checkpoint, the shape in current model is torch.Size([105, 1024]).
    size mismatch for classifier.bias: copying a param with shape torch.Size([59]) from checkpoint, the shape in current model is torch.Size([105]).

I see that the problem probably is a mismatch from labels size between both Datasets. I've searched a bit and I've found a recommendation to use ignore_mismatched_sizes=True as and argument for pretrained. But I keep receiving the same error.

Here is part of my code when trying to predict on unseen data:

from transformers import BertForSequenceClassification

# Just right before the actual usage select your hardware
device = torch.device('cuda') # or cpu
model = model.to(device)      # send your model to your hardware



model = BertForSequenceClassification.from_pretrained("neuralmind/bert-large-portuguese-cased",
                                                      num_labels=len(label_dict),
                                                      output_attentions=False,
                                                      output_hidden_states=False,
                                                      ignore_mismatched_sizes=True)

model.to(device)

model.load_state_dict(torch.load('finetuned_BERT_epoch_2_full-Copy1.model', map_location=torch.device('cuda')))

_, predictions, true_vals = evaluate(dataloader_validation)
accuracy_per_class(predictions, true_vals)

Could someone help me how could I deal with it? I don't know what more can I do!

Any help I'm very grateful!

1 Answers

Your new dataset has 105 classes while your model was trained for 59 classes. As you have already mentioned, you can use ignore_mismatched_sizes to load your model. This parameter will load the the embedding and encoding layers of your model, but will randomly initialize the classification head:

model = BertForSequenceClassification.from_pretrained("finetuned_BERT_epoch_2_full-Copy1.model",
                                                      num_labels=105,
                                                      output_attentions=False,
                                                      output_hidden_states=False,
                                                      ignore_mismatched_sizes=True)

In case you want to keep the classification layer of the 59 labels and add 46 labels, you can refer to this answer. Please also note the comments of this answer, because this approach does not provide any meaningful results due to the random initialization for the new labels.

Related