sklearn
compute_class_weight
large dataset
machine learning
Python

sklearn utils compute_class_weight function for large dataset

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

sklearn.utils.class_weight.compute_class_weight is a small helper, but questions about it usually come from a larger problem: class imbalance on a dataset big enough that naive preprocessing becomes expensive. The important part is not the helper itself. It is understanding the weight formula and computing class counts efficiently.

For the common balanced mode, the weight for each class is based on the total number of samples divided by the number of classes and by the count of that class. Once you know that, you can decide whether the built-in helper is enough or whether a custom counting path fits the dataset better.

How compute_class_weight Works

For a normal in-memory label array, the helper is straightforward:

python
1import numpy as np
2from sklearn.utils.class_weight import compute_class_weight
3
4y = np.array([0, 0, 0, 1, 1, 2])
5classes = np.unique(y)
6
7weights = compute_class_weight(
8    class_weight="balanced",
9    classes=classes,
10    y=y,
11)
12
13print(dict(zip(classes, weights)))

This is appropriate when y already fits comfortably in memory and class labels are easy to enumerate.

Why Large Datasets Need a Different Approach

On a large dataset, the expensive part is rarely the weight formula. The real cost is usually one of these:

  • loading labels into memory all at once
  • converting labels repeatedly between formats
  • counting labels inefficiently inside a preprocessing loop
  • computing weights on one distribution and training on a different filtered distribution

That changes the design question. Instead of asking whether compute_class_weight scales, ask how to derive stable counts from the actual training labels you will use.

Fast Counting for Integer Labels

If labels are dense nonnegative integers, numpy.bincount is usually the fastest simple option:

python
1import numpy as np
2
3y = np.array([0, 0, 0, 1, 1, 2], dtype=np.int64)
4counts = np.bincount(y)
5
6total = y.size
7n_classes = np.count_nonzero(counts)
8
9weights = {
10    cls: total / (n_classes * count)
11    for cls, count in enumerate(counts)
12    if count > 0
13}
14
15print(weights)

This does the same job as the balanced helper for many classification pipelines and is easy to reason about.

Streaming Counts for Very Large Data

If the label vector does not fit in memory, count in chunks and compute weights after the pass is complete. The formula only needs final class counts.

python
1from collections import Counter
2
3
4def balanced_weights(counts):
5    total = sum(counts.values())
6    n_classes = len(counts)
7    return {
8        label: total / (n_classes * count)
9        for label, count in counts.items()
10    }
11
12
13counter = Counter()
14for chunk in ([0, 0, 1], [1, 2, 2], [2, 2, 2]):
15    counter.update(chunk)
16
17print(counter)
18print(balanced_weights(counter))

The same strategy works if labels come from chunked CSV reads, parquet batches, or a streaming data source.

Passing Weights into Training

Be careful about the interface expected by the model library. Scikit-learn estimators often accept a class_weight mapping keyed by class label, while some APIs use per-sample weights instead. Those are related but not interchangeable.

For example, after computing a mapping such as 0: 0.66, 1: 1.0, 2: 2.0, you still need to confirm that the estimator consumes it in the format you expect. If a training job uses only a sampled subset of the data, recalculate weights on that exact subset.

Common Pitfalls

  • Treating compute_class_weight as the performance bottleneck when label loading is the real issue.
  • Using numpy.bincount on string labels or sparse integer label spaces without remapping first.
  • Computing weights before train-validation splitting and then applying them to a different distribution.
  • Confusing class weights with sample weights.
  • Assuming heavier balancing always improves the validation metric.

Summary

  • 'compute_class_weight is a convenience wrapper around a simple class-frequency formula.'
  • For large datasets, efficient counting matters more than the helper call itself.
  • Use numpy.bincount for dense integer labels and chunked counting for out-of-memory data.
  • Compute weights from the exact training distribution you will actually fit.
  • Verify whether the target estimator expects class weights or per-sample weights.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.