What loss function to use for imbalanced classes (using PyTorch)?

Question

I have a dataset with 3 classes with the following items:

• Class 1: 900 elements
• Class 2: 15000 elements
• Class 3: 800 elements

I need to predict class 1 and class 3, which signal important deviations from the norm. Class 2 is the default “normal” case which I don’t care about.

What kind of loss function would I use here? I was thinking of using CrossEntropyLoss, but since there is a class imbalance, this would need to be weighted I suppose? How does that work in practice? Like this (using PyTorch)?

summed = 900 + 15000 + 800
weight = torch.tensor([900, 15000, 800]) / summed
crit = nn.CrossEntropyLoss(weight=weight)

Or should the weight be inverted? i.e. 1 / weight?

Is this the right approach to begin with or are there other / better methods I could use?

Thanks

Answer 1

What kind of loss function would I use here?

Cross-entropy is the go-to loss function for classification tasks, either balanced or imbalanced. It is the first choice when no preference is built from domain knowledge yet.

This would need to be weighted I suppose? How does that work in practice?

Yes. Weight of class $c$ is the size of largest class divided by the size of class $c$.

For example, If class 1 has 900, class 2 has 15000, and class 3 has 800 samples, then their weights would be 16.67, 1.0, and 18.75 respectively.

You can also use the smallest class as nominator, which gives 0.889, 0.053, and 1.0 respectively. This is only a re-scaling, the relative weights are the same.

Is this the right approach to begin with or are there other / better methods I could use?

Yes, this is the right approach.

EDIT:

Thanks to @Muppet, we can also use class over-sampling, which is equivalent to using class weights. This is accomplished by WeightedRandomSampler in PyTorch, using the same aforementioned weights.

Answer 2

I am totaly agree with @Esmailian

def compute_pos_weights(cls_repr: torch.Tensor) -> torch.Tensor:
    total_weight = cls_repr.sum()
    weights = 1/torch.div(cls_repr, total_weight)
    # Standardize the weights
    return torch.div(weights, torch.min(weights))