Train-Test Split & Cross-Validation
train_test_split() as a black box in earlier lessons — this lesson opens it up. You'll see why a single random split can give a misleadingly good or bad performance estimate, how k-fold cross-validation solves that by testing on every row exactly once, and how stratified splitting keeps class proportions intact for classification problems.
train_test_split() in Detail
train_test_split, from sklearn.model_selection, randomly shuffles the rows of X and y together and divides them into a training portion and a held-out test portion.
from sklearn.model_selection import train_test_split X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.25, # 25% of rows held out for testing random_state=42 # fixes the shuffle so results are reproducible ) print("Train rows:", len(X_train), "| Test rows:", len(X_test))
Why a Single Split Can Be Noisy
A single train_test_split() call picks one particular random sample of rows for testing. On a small or unevenly-distributed dataset, that one sample might happen to be unusually easy or unusually hard — purely by chance. Change random_state from 42 to 7, and the reported test accuracy can shift meaningfully, even though nothing about the model or the data actually changed.
The fix isn't to abandon the train/test split — it's to average performance over MULTIPLE different splits, so a single unlucky sample can't dominate the result. That's exactly what cross-validation does.
K-Fold Cross-Validation
K-fold cross-validation splits the data into k equally-sized "folds." It then runs k separate rounds: in each round, one fold is held out as the test set and the model trains on the remaining k − 1 folds. Every row gets used as test data exactly once, across the k rounds — so the final average score isn't at the mercy of one lucky or unlucky split.
| Fold 1 | Fold 2 | Fold 3 | Fold 4 | Fold 5 | |
|---|---|---|---|---|---|
| Round 1 | TEST | train | train | train | train |
| Round 2 | train | TEST | train | train | train |
| Round 3 | train | train | TEST | train | train |
| Round 4 | train | train | train | TEST | train |
| Round 5 | train | train | train | train | TEST |
5-fold cross-validation: each fold is the test set exactly once, across 5 training rounds.
from sklearn.model_selection import cross_val_score from sklearn.tree import DecisionTreeClassifier model = DecisionTreeClassifier(random_state=42) # cv=5 runs 5-fold cross-validation and returns one score per fold scores = cross_val_score(model, X, y, cv=5) print("Fold scores:", scores) print("Mean accuracy:", scores.mean()) print("Std deviation:", scores.std())
Note that cross_val_score handles the splitting internally — you pass the FULL X and y, not a pre-split train/test pair. A common workflow is to use cross-validation on the training set for model comparison and tuning, then do one final check on a completely separate test set you held out from the start.
Stratified Splitting for Classification
A plain random split can, by chance, put too many or too few examples of a rare class into the test set — especially with imbalanced classification data (say, 90% "no churn" and 10% "churn"). Stratified splitting forces every split to preserve the original class proportions.
from sklearn.model_selection import train_test_split # stratify=y keeps the class balance of y the same in both splits X_train, X_test, y_train, y_test = train_test_split( X, y, test_size=0.2, random_state=42, stratify=y )
The same idea extends to cross-validation with StratifiedKFold, which builds folds that each preserve the overall class ratio — the recommended default for classification cross-validation.
from sklearn.model_selection import StratifiedKFold, cross_val_score from sklearn.tree import DecisionTreeClassifier model = DecisionTreeClassifier(random_state=42) skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42) scores = cross_val_score(model, X, y, cv=skf) print("Stratified fold scores:", scores) print("Mean accuracy:", scores.mean())
stratify works on discrete class labels, so it's a classification-only tool. For regression, plain KFold (no stratification) is the standard choice, since the target is continuous rather than a set of categories.Lesson Summary
cross_val_score) tests on every row exactly once across k rounds, giving a more reliable estimate.test_size=0.2 mean in train_test_split?stratify=y in train_test_split guarantee for a classification problem?Get hands-on with all three tools from this lesson.
Run
train_test_split(X, y, test_size=0.2, random_state=1) and then again with random_state=99. Train the same model on each split and compare the two test accuracies. Are they identical? Write a sentence on why or why not.
Using
cross_val_score(model, X, y, cv=5) from Section 3, print the 5 fold scores, the mean, and the standard deviation. Compare the mean to the single-split score from Task 1.
Split an imbalanced classification dataset both with and without
stratify=y. Using pandas.Series(y_test).value_counts(normalize=True) on each result, confirm the stratified version's class proportions match the full dataset's more closely.
💡 Show hints if you're stuck
- Task 1: The two accuracies will usually differ at least slightly — different rows end up in the test set each time, even though the model and data are unchanged.
- Task 2: The cross-validation mean is generally the more trustworthy number since it isn't dependent on one particular split.
- Task 3: Without stratification, a rare class's proportion in the test set can swing noticeably from run to run — with
stratify=y, it should stay very close to the original ratio every time.