Custom BCEWithLogitsLoss for Sequence Classification using Auto Model

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.

1 Like

Hmm…

BCEWithLogitsLoss is for Binary 0 or 1 use torch.nn.CrossEntropyLoss