Handling Imbalanced Classes with Weighted Cross-Entropy

1. The Peril of Imbalanced Data

When training a neural network, the default objective is to minimize the overall loss, which often correlates with maximizing accuracy. However, when one class vastly outnumbers others, a "lazy" model can achieve high accuracy by simply always predicting the majority class.

Imagine your dataset has 80 'Adequate' images, 15 'Good' images, and only 5 'Poor' images. A model that learns to classify every image as 'Adequate' would achieve 80% accuracy, a seemingly impressive score. Yet, it would be completely useless for your project's goal of identifying violations, as it would never detect a 'Poor' quality meal.

This scenario highlights why raw accuracy is a deeply flawed metric for imbalanced problems. The Paperspace article "Class Imbalance in Image Datasets" provides an excellent demonstration of this issue.

2. Strategies to Address Imbalance

Now that we understand the problem, let's look at the solutions. There are two main families of techniques:

1. Data-Level Methods: These involve resampling the dataset to create a balanced distribution before training.

   • Undersampling: Randomly removes samples from the majority class. This is simple but can discard potentially useful information.
   • Oversampling: Randomly duplicates samples from the minority class. This ensures no data is lost but can lead to the model overfitting on the repeated examples.

2. Algorithm-Level Methods: These modify the training algorithm itself to account for the imbalance. The most common technique, and our focus for this lesson, is class weighting.

Instead of changing the data, class weighting changes the loss function. We assign a higher penalty for misclassifying an example from a minority class. This effectively tells the model, "Making a mistake on a 'Poor' meal is much worse than making a mistake on an 'Adequate' meal, so pay extra attention to getting the 'Poor' ones right."

3. Implementing Weighted Cross-Entropy Loss

In the previous lesson, you defined your loss function like this:
criterion = nn.CrossEntropyLoss()

PyTorch makes it remarkably simple to add class weights. The CrossEntropyLoss function accepts an optional weight argument, which is a tensor containing the weight for each class.

The key is to determine what those weights should be. A standard approach is to use the inverse of the class frequency.

Calculating Class Weights

First, you need to count the number of samples in each of your classes (Good, Adequate, Poor). Let's say your counts are:

• n_good = 150
• n_adequate = 750
• n_poor = 100

The weight for a class c is typically calculated as the total number of samples divided by the number of samples in that class. However, a simpler and more common formula for inverse frequency is:

Using this formula with our example counts:

• weight_good = 750 / 150 = 5.0
• weight_adequate = 750 / 750 = 1.0
• weight_poor = 750 / 100 = 7.5

This calculation assigns the lowest weight to the majority class ('Adequate') and the highest weight to the most underrepresented class ('Poor').

Applying Weights in PyTorch

Once you have your weights, you create a tensor and pass it to the loss function. It's crucial to move this tensor to the same device as your model and data.

# Assume 'device' is already defined (e.g., 'cuda')
# Class counts for 'Good', 'Adequate', 'Poor'
# You would calculate these from your actual dataset
class_counts = [150, 750, 100]
max_count = max(class_counts)




# Calculate weights
weights = [max_count / count for count in class_counts]




# Create a tensor and move it to the GPU
class_weights = torch.tensor(weights, dtype=torch.float32).to(device)




# Define the loss function with weights
criterion = nn.CrossEntropyLoss(weight=class_weights)




# The rest of your training loop remains the same!

This single change forces your optimizer to work harder to classify 'Good' and especially 'Poor' meals correctly.

The following resources provide excellent code examples and further explanation of this process. The Medium article gives a very clear formula for calculating weights, and the video from Aladdin Persson shows the implementation in just a few lines of code.

For a more in-depth, end-to-end example, the Paperspace article you reviewed earlier also contains the full implementation. It modifies the training loop to track recall and precision, showing the direct impact of applying class weights.

4. An Alternative: Balanced Sampling

While weighted loss is a very effective and easy-to-implement strategy, it's worth knowing about the primary data-level alternative: oversampling using PyTorch's WeightedRandomSampler.

Instead of weighting the loss, this technique adjusts the DataLoader. It creates a sampler that picks data points for each batch according to a specific weight, ensuring that on average, each batch contains a balanced representation of all classes. The model sees more examples of the minority class, effectively balancing the data on-the-fly.

The Aladdin Persson video provides a clear tutorial on this method if you find that weighted loss is not sufficient for your needs.

Conclusion

You have now tackled a fundamental challenge in applied machine learning. By understanding and addressing class imbalance, you are moving beyond a naive implementation and building a model that is far more likely to be useful for your actual project goal.

Here are your key takeaways:

• Class imbalance is common in real-world datasets and can lead to biased models.
• Accuracy is a misleading metric for imbalanced problems; a model can achieve a high score by ignoring the minority classes you care about.
• Weighted loss functions penalize errors on minority classes more heavily, forcing the model to learn their features.
• Implementing weighted CrossEntropyLoss in PyTorch is a straightforward and effective way to handle imbalance by passing a weight tensor during initialization.

Now that you have a strategy to train a more "fair" model, the next crucial step is to properly measure its performance. In our next lesson, we will dive deep into evaluating your classifier. You will learn to implement and interpret essential metrics like precision, recall, F1-score, and the confusion matrix to gain a true understanding of your model's strengths and weaknesses.