TopolearnA model scored on its own training data looks far better than it really is.
The goal of a model is to do well on data it has not seen. Measuring it on the training data rewards memorising, so you hold out a test set and score on that.
train_test_split(X, y, test_size=0.2, random_state=0) shuffles the rows and returns four pieces in a fixed order: X_train, X_test, y_train, y_test. Useful arguments:
test_size/train_size: a fraction (0.2) or an absolute count.random_state: fixes the shuffle so results are reproducible.stratify=y: keeps class proportions equal in both parts, important for imbalanced or small classification datasets.shuffle=False: keep the original order, for time-ordered data where the test set must come after the training set.
You can pass more than two arrays (for example sample weights or IDs) and each is split the same way.
Treat the test set as sealed. Fit every learned step, including scalers and encoders, on the training part only, and do all model choices with the training data (cross-validation helps here). Look at the test score once, at the end. If you keep tweaking until the test score improves, it stops being an honest estimate.
A single split is noisy on small datasets: a different random_state can change the score noticeably. Cross-validation, covered later, averages over several splits.
Example
import numpy as np
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
X, y = make_classification(n_samples=1000, weights=[0.8, 0.2], random_state=0)
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.25, stratify=y, random_state=42)
print(X_train.shape, X_test.shape) # (750, 20) (250, 20)
print(y_train.mean().round(3), y_test.mean().round(3)) # 0.203 0.204 (same class balance)
idx = np.arange(10)
print(train_test_split(idx, shuffle=False, test_size=3)) # [array([0, 1, 2, 3, 4, 5, 6]), array([7, 8, 9])]Remember
- Return order is X_train, X_test, y_train, y_test
- Use
stratify=yfor classification andrandom_statefor reproducibility - Look at the test score once, after all decisions are made
Watch out for
- Fitting a scaler or encoder on the full data before splitting, which leaks test information.