ML Train/Test Split

To know if a model really works, you must test it on data it has never seen.

Why Split the Data?

Imagine a teacher gives students the exact exam questions to practise. Everyone gets full marks, but nobody has really learned. The marks tell you nothing.

Models are the same. If you test a model on the data it trained on, it can simply remember the answers. So we keep some data aside:

The test set is the "final exam". The model must never see it during training.

Using train_test_split

scikit-learn has a ready-made function for this:

from sklearn.model_selection import train_test_split

X = [[i] for i in range(10)]          # 10 examples, 1 feature
y = [0, 0, 0, 0, 0, 1, 1, 1, 1, 1]

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42
)

print("Training rows:", len(X_train))
print("Test rows:", len(X_test))
print("Test data:", X_test, y_test)

Training rows: 8
Test rows: 2
Test data: [[8], [1]] [1, 0]

The function shuffles the rows and returns four parts, always in this order: X_train, X_test, y_train, y_test.

Important parameters

A Full Example: Train, Then Test

This example uses the Iris dataset, which comes with scikit-learn. It has measurements of 150 flowers from 3 species.

from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

X, y = load_iris(return_X_y=True)

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=1, stratify=y
)

model = KNeighborsClassifier(n_neighbors=5)
model.fit(X_train, y_train)          # learn only from training data

print("Test accuracy:", round(model.score(X_test, y_test), 2))

Test accuracy: 0.95

The model was right on 95% of flowers it had never seen. That number is a fair estimate of how it will do on new flowers.

Keeping Classes Balanced with stratify

If your data has 3 classes in equal numbers, a random split might put too many of one class in the test set. stratify=y prevents this:

import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split

X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=0, stratify=y
)

print("Whole data:", np.bincount(y))
print("Test set:  ", np.bincount(y_test))

Whole data: [50 50 50]
Test set:   [10 10 10]

Stratifying matters most when one class is rare, such as fraud cases that are only 1% of transactions.

Common Mistakes

• Preparing data before splitting. If you fit a scaler on all the data, information from the test set leaks into training. Split first, then fit the scaler on the training set only.
• Testing again and again on the same test set while you tune your model. You slowly "fit" the test set. Use cross-validation for tuning (see Overfitting).
• Shuffling time data. For sales or weather over time, train on the past and test on the future, not a random mix.

Exercise