Data splitting

Data is split into train and test sets with train_test_split.

stratify

If classes are imbalanced, a random split may leave few or none of one class in the test set. stratify=y keeps the class proportions in both sets:

X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, stratify=y, random_state=0)

In classification problems stratify=y is almost always a good default.

Validation set

If you choose model settings (hyperparameters) by looking at the test set, the test set is indirectly "learned" too. That's why a three-way split is used:

X_tmp, X_test, y_tmp, y_test = train_test_split(X, y, test_size=0.2, random_state=0)
X_train, X_val, y_train, y_val = train_test_split(X_tmp, y_tmp, test_size=0.25, random_state=0)
# result: 60% train, 20% validation, 20% test

Settings are chosen with the validation set; the test set is used only once, at the very end.

Exercise Stratified split

Split using y made of 16 zeros and 4 ones with test_size=0.25, stratify=y and random_state=0. Print the sum of the test labels (the number of 1s).

Show example solution

from sklearn.model_selection import train_test_split

X = [[i] for i in range(20)]
y = [0] * 16 + [1] * 4

X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, stratify=y, random_state=0)
print(sum(y_test))

Test yourself

With labels that are 90% 0 and 10% 1, what is the share of 1s in the test set when <code>stratify=y</code> is used?
- About 10%
- 50%
- 0%
- Random
Show answer

About 10% — Stratified splitting keeps the original class proportion in each set.

Keep the class proportions in the split.

train_test_split(X, y, test_size=0.2, ____=y)

Show answer

train_test_split(X, y, test_size=0.2, stratify=y)

Match the sets with their roles.

Training

The model learns
Validation

Settings are chosen
Test

Measured once at the end
stratify

Keeps class proportion

Next lesson: Summary - Data preprocessing

We Are Enhancing Your Experience

Cookie Preferences

Strictly Necessary Cookies

Analytical Cookies