Class weights in loss function for imbalanced dataset - how to work with an imbalanced dataset?

[upasana-crypto]

Did someone work with imbalanced dataset and used class weights in the loss function ? Can you share example of train and validation because we need to pass weights in the training but not in validation. I am a bit confused on how to do that.

[mrdbourke]

Hi @upasana-crypto,

You can use the weight parameter as part of the torch.nn.CrossEntropyLoss loss function.

This can be used for training and validation.

For example:

import torch
import torch.nn as nn

class ImbalancedModel(nn.Module):
    def __init__(self):
        super().__init__()
        self.model = nn.Linear(10, 2)
        
    def forward(self, x):
        return self.model(x)

# Define the model and the optimizer
model = ImbalancedModel()
optimizer = torch.optim.Adam(model.parameters())

# Define the criterion with class-specific weights
loss_fn = nn.CrossEntropyLoss(weight=torch.tensor([0.5, 2.0]))

# Create datasets/dataloaders
...

# Train the model
for inputs, labels in train_dataloader:
    # Compute the model output and the loss
    output = model(inputs)
    loss = loss_fn(output, labels)

    # Backpropagate the loss and update the model's parameters
    loss.backward()
    optimizer.step()

# Validation code (can use the same loss function)
...

In this example, the ImbalancedModel class defines a linear model that takes 10-dimensional inputs and produces 2-dimensional outputs.

The CrossEntropyLoss (loss_fn) is defined with class-specific weights, where the weight for the first class is 0.5 and the weight for the second class is 2.0.

During training, the loss is computed using these class-specific weights, which means that the model will pay more attention to the second class (which has a higher weight) and less attention to the first class (which has a lower weight).

This can help the model learn to better distinguish between the two classes, even if they are imbalanced in the training data.

Same loss function for validation and testing as training

You should use the same class-specific weights during validation and testing as you used during training.

If you use different weights during validation or testing, the model's performance will be evaluated using a different loss function, which could lead to misleading or inaccurate results.