Train Test Split
Python · Machine Learning (Scikit-Learn)
Theory
A model must be judged on data it has never seen, otherwise you only measure how well it memorised. train_test_split from sklearn.model_selection shuffles the rows and divides X and y into a training set (used to fit) and a test set (held back for evaluation). test_size=0.25 keeps 25% for testing, and it returns four objects in a fixed order: X_train, X_test, y_train, y_test. random_state fixes the shuffle so the split is the same every run.
For classification, a plain random split can leave the test set with too few or too many of a rare class. stratify=y keeps the class proportions the same in both parts: with 4 frauds in 20 rows (20%), the stratified test set of 5 rows gets exactly 1 fraud. Fit scalers, encoders and models on the training part only, and for time-ordered data split by date instead of shuffling.
Why we use it
It gives an honest estimate of performance on new data and exposes overfitting. Stratifying keeps rare classes represented, so precision and recall on the test set mean something.
When to use it
Split once at the start, before any fitting or scaling. Add stratify=y for classification, especially with imbalanced classes. For small datasets prefer cross-validation, and for time series use a date cut-off or TimeSeriesSplit.
How it works
- Features X and target y
- Classification with imbalanced classes? Yes: Add stratify=y. No: Plain random split.
- Shuffle with a fixed random_state
- Hold back test_size of the rows as the test set
- Fit everything on X_train, y_train only
- Evaluate once on X_test, y_test
Syntax
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y)Example
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42, stratify=y)
print(y_train.mean(), y_test.mean())Key points
- Returns X_train, X_test, y_train, y_test in that order
- random_state makes the split reproducible
- stratify=y keeps class proportions in both parts
- Fit preprocessing on the training set only to avoid leakage
- Time series: split by date, do not shuffle