lightly-ai / lightly-ai/lightly-train

[FEAT] Support class weights in the loss to handle imbalanced datasets

Open
#965 0 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement tracking issue
Dominant language
Python
Stars
1.7k
Forks
116
Avg merge
2d 21h
Merged PRs (30d)
6

Description

## Description

Many real datasets are imbalanced: some classes have far more images than others. The model then sees the common classes much more often, and usually ends up good at those and bad at the rare ones.

A common way to counter this is *class weighting*: we tell the loss function that a mistake on a rare class counts more than a mistake on a common class. PyTorch supports this directly (`weight` for `CrossEntropyLoss`, `pos_weight` for `BCEWithLogitsLoss`), but LightlyTrain does not expose it for any task today. Would e.g. be helpful in #964.

We would like to add one consistent way to set class weights, and roll it out task by task.

## Expected behavior

Every task gets a `class_weights` training argument with the same meaning:

- `None` (default): no class weighting, exactly like today.
- `"auto"`: LightlyTrain counts the classes in the training data and gives rare classes a higher weight.
- A dictionary from class name to weight, for users who want full control, for example `{"cat": 1.0, "dog": 3.5}`.

This reuses the `"auto"` pattern that LightlyTrain already uses for arguments like `weight_decay` and `lr_warmup_steps`.

Rules that should hold for every task:

- **The default stays `None`.** Class weighting is a trade-off: it usually improves accuracy on rare classes and lowers overall accuracy. Users should ask for it, and existing runs should not change.
- **Weights are keyed by class name, not by position in a list.** Classes can be dropped with `ignore_classes` and are renumbered internally, so a plain list would silently attach the wrong weight to the wrong class.
- **The formula behind `"auto"` is decided per task.** Inverse class frequency is a good fit for classification, but segmentation counts pixels instead of images and normally needs a gentler formula.

## Sub-issues

- [ ] Image classification

Sub-issues for the other tasks (semantic segmentation, object detection) will be added later. If you would like to work on one, please leave a comment first so we do not duplicate work.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by tracing the image-classification task's training arguments and loss construction, then identify how class names, ignored classes, and training-data counts are represented. Define the shared None, "auto", and name-to-weight behavior for image classification, while preserving current defaults; done means the classification sub-issue works with tests covering all three input forms and ignored classes.

Written by the indexing model from the issue text.

Assessment

Tech stack
python, pytorch
Domain
computer-vision, machine-learning
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
45/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.