Train/Test Split
Hold back part of the data to honestly measure how the model performs on unseen inputs.
Syntax
train_test_split(X, y, test_size=0.2, random_state=42)If you test a model on the same data it learned from, it looks great but tells you nothing about new data. So we split the data.
How it works
train_test_split randomly divides rows into a training set (to learn from) and a test set (to judge on). A common split is 80/20.
Example
from sklearn.model_selection import train_test_split
import numpy as np
X = np.arange(20).reshape(10, 2) # 10 samples, 2 features
y = np.array([0,1,0,1,0,1,0,1,0,1])
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42)
print(X_train.shape) # (8, 2)
print(X_test.shape) # (2, 2)When to use it
- An ML engineer reserves 20% of a labelled dataset as a test set before any exploratory analysis to guarantee a completely unseen evaluation benchmark.
- A researcher uses stratified splitting on an imbalanced medical dataset to ensure the rare positive class appears in both train and test sets proportionally.
- A data scientist creates train, validation, and test splits by applying train_test_split twice to tune hyperparameters without touching the test set.
More examples
Basic 80/20 train-test split
Splits features and labels into 80% train and 20% test subsets with a fixed random seed for reproducibility.
from sklearn.model_selection import train_test_split
import pandas as pd
df = pd.read_csv('churn.csv')
X = df.drop(columns=['churned'])
y = df['churned']
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42)
print(X_train.shape, X_test.shape)Stratified split for imbalanced data
Uses stratify=y so the rare class rate (5%) is preserved in both train and test sets, preventing the test set from being all-negative.
from sklearn.model_selection import train_test_split
import numpy as np
X = np.random.randn(1000, 5)
y = np.array([0]*950 + [1]*50) # 5% positive class
X_tr, X_te, y_tr, y_te = train_test_split(
X, y, test_size=0.2, stratify=y, random_state=0)
print('Train positive rate:', y_tr.mean().round(3))
print('Test positive rate:', y_te.mean().round(3))Three-way train / val / test split
Applies train_test_split twice to carve out separate validation and test sets, keeping the test set truly unseen until final evaluation.
from sklearn.model_selection import train_test_split
from sklearn.datasets import load_breast_cancer
X, y = load_breast_cancer(return_X_y=True)
X_temp, X_test, y_temp, y_test = train_test_split(X, y, test_size=0.15, random_state=1)
X_train, X_val, y_train, y_val = train_test_split(X_temp, y_temp, test_size=0.18, random_state=1)
print(X_train.shape, X_val.shape, X_test.shape)
Discussion