undersampling
machine learning
train/test split
data preprocessing
class imbalance

Undersampling before or after Train/Test Split

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Undersampling should be applied after the train/test split, not before it. More precisely, it should be applied only to the training data, and if you are using cross-validation, it should happen separately inside each training fold.

The reason is data leakage. If you undersample before splitting, information from the full dataset influences what ends up in training and test data, and the test distribution stops representing the real problem you actually care about.

Why Splitting Comes First

The test set is supposed to simulate unseen production data. That means it should preserve the natural class imbalance of the original dataset unless you have a very unusual evaluation goal.

If you undersample first, you are changing the population before creating the evaluation set. That makes the test set easier and less realistic, which can produce misleading performance numbers.

So the safe order is:

  1. split the data
  2. undersample only the training portion
  3. train the model on the resampled training data
  4. evaluate on the untouched test data

Correct Workflow with scikit-learn and imbalanced-learn

Here is a simple example using RandomUnderSampler.

python
1from sklearn.model_selection import train_test_split
2from imblearn.under_sampling import RandomUnderSampler
3from collections import Counter
4
5X = [[i] for i in range(100)]
6y = [0] * 90 + [1] * 10
7
8X_train, X_test, y_train, y_test = train_test_split(
9    X, y, test_size=0.2, stratify=y, random_state=42
10)
11
12print("Before undersampling:", Counter(y_train))
13
14sampler = RandomUnderSampler(random_state=42)
15X_train_res, y_train_res = sampler.fit_resample(X_train, y_train)
16
17print("After undersampling:", Counter(y_train_res))
18print("Untouched test set:", Counter(y_test))

The training set becomes balanced, but the test set keeps the original imbalance. That is exactly what you want for honest evaluation.

Why Undersampling Before the Split Is Wrong

This is the pattern to avoid:

python
1sampler = RandomUnderSampler(random_state=42)
2X_res, y_res = sampler.fit_resample(X, y)
3
4X_train, X_test, y_train, y_test = train_test_split(
5    X_res, y_res, test_size=0.2, random_state=42
6)

This leaks information from the full dataset into the resampling step and produces a test set whose class balance no longer matches the original problem.

The result is often an overly optimistic estimate of model quality.

Cross-Validation Needs the Same Discipline

The same rule applies when using cross-validation. Resampling should happen inside each training fold, not once on the entire dataset before the folds are created.

A good way to do that is with an imblearn pipeline.

python
1from imblearn.pipeline import Pipeline
2from imblearn.under_sampling import RandomUnderSampler
3from sklearn.linear_model import LogisticRegression
4from sklearn.model_selection import cross_val_score, StratifiedKFold
5
6pipeline = Pipeline([
7    ("undersample", RandomUnderSampler(random_state=42)),
8    ("model", LogisticRegression(max_iter=1000))
9])
10
11cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
12scores = cross_val_score(pipeline, X, y, cv=cv, scoring="f1")
13
14print(scores)

This ensures the sampler sees only the training portion of each fold.

Keep Evaluation Metrics Appropriate

Because the test set remains imbalanced, plain accuracy is often a weak metric. In class-imbalance problems, metrics such as precision, recall, F1, PR AUC, or ROC AUC are usually more informative.

That is not a reason to rebalance the test set. It is a reason to choose metrics that reflect the real problem.

Common Pitfalls

  • Undersampling before the train/test split and leaking information into evaluation.
  • Balancing the test set and then believing the measured performance reflects production behavior.
  • Applying the sampler once before cross-validation instead of inside each fold.
  • Using only accuracy on a strongly imbalanced problem.
  • Forgetting stratification during the original split and accidentally distorting the class ratios even before resampling.

Summary

  • Undersample after the train/test split, not before.
  • Apply undersampling only to the training data.
  • Keep the test set untouched so it reflects the real class distribution.
  • In cross-validation, resample inside each training fold through a pipeline.
  • Use evaluation metrics suited to imbalanced classification rather than trying to make the test set artificially balanced.

Course illustration
Course illustration

All Rights Reserved.