I am training a BERT based model with AutoModelForSequenceClassification
. I want to use class weights for the loss function so I followed this discussion and implemented a BCEWithLogitsLoss
.
Problem is, HuggingFace advices using num_labels=2
for Binary Classification with the Auto Class. When I try to implement BCEWithLogitsLoss
I run into the error that
Value Error: Target size (torch.Size([32])) must be the same as input size (torch.Size([32, 2]))
I understand why this is happening, the size of the b_labels
is the ground truth while the logits
have the shape [32, 2]
but I can’t figure out how to fix it. Here’s the code:
criterion = nn.BCEWithLogitsLoss(weight=class_weights, reduction='mean')
logits = outputs.logits
loss = criterion(logits, b_labels)
Here, class_weights is tensor([1.0712, 0.9377])
Should I just use CrossEntropyLoss
instead? I have dataset imbalance that’s why I want to switch to custom loss and see what happens. Please guide me towards a solution.