Train/Test Split

A model is useful only if it works on data it has never seen. To measure that, you hide part of your data during training and use it afterwards as an exam. This is the train/test split, and getting it right is the difference between an honest score and a score that fools you.

In this lesson you will see why testing on training data misleads, use train_test_split with all its important options, build a validation set, handle time-ordered data, and learn to spot data leakage.

Why you need a test set

Imagine a student who memorises the answers to last year's paper. They score 100% on that paper and fail the real exam. A flexible model can do the same thing: memorise the training rows instead of learning the general pattern. Here an unlimited decision tree is trained on breast cancer biopsies and checked on both sets.

from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=42)

tree = DecisionTreeClassifier(random_state=42).fit(X_train, y_train)
print(f"Accuracy on training data: {tree.score(X_train, y_train):.3f}")
print(f"Accuracy on test data:     {tree.score(X_test, y_test):.3f}")

Accuracy on training data: 1.000
Accuracy on test data:     0.951

The perfect training score tells you nothing; the tree has simply memorised every training row. The test score is the honest estimate of how it will do on new patients. The gap between the two numbers is a warning sign called overfitting, which you will study in Overfitting & the Bias-Variance Trade-off.

How train_test_split works

train_test_split shuffles the rows and cuts them into two parts. You pass it your arrays (features, target, or both), and it returns each array split in two, always in the order train, test for every array you passed:

• test_size is the fraction (or number) of rows for the test set. Values from 0.2 to 0.3 are common.
• random_state fixes the shuffle so you get the same split every run. Without it, every run gives a different split and different scores.
• shuffle=True is the default; you will see when to turn it off below.

import numpy as np
from sklearn.model_selection import train_test_split

roll_numbers = np.arange(101, 111)                  # ten students
passed = np.array([1, 1, 0, 1, 0, 1, 1, 0, 1, 1])

r_train, r_test, p_train, p_test = train_test_split(
    roll_numbers, passed, test_size=0.3, random_state=42)
print("Train rolls:", r_train)
print("Test rolls: ", r_test, "labels:", p_test)

print("Same seed:  ", train_test_split(roll_numbers, test_size=0.3, random_state=42)[1])
print("Other seed: ", train_test_split(roll_numbers, test_size=0.3, random_state=7)[1])

Train rolls: [101 108 103 110 105 104 107]
Test rolls:  [109 102 106] labels: [1 1 1]
Same seed:   [109 102 106]
Other seed:  [109 106 101]

The rows stay matched: roll number 109 and its label go to the same side. The same random_state always gives the same split, which makes your results reproducible. The exact number you choose does not matter; this course uses 42 out of habit.

Keep class proportions with stratify

Notice that all three test labels above are 1. With small or unbalanced data, a random cut can easily give a test set that looks nothing like the full data. Suppose 10% of 100 UPI payments are fraud. A test set of 20 should contain about 2 fraud cases, but plain shuffling might give 0 or 5. stratify=y makes the split keep the class proportions.

import numpy as np
from sklearn.model_selection import train_test_split

y = np.array([0] * 90 + [1] * 10)   # 10 fraud payments in 100
X = np.arange(100).reshape(-1, 1)

for seed in [0, 1, 2, 3]:
    *_, plain = train_test_split(X, y, test_size=0.2, random_state=seed)
    *_, strat = train_test_split(X, y, test_size=0.2, random_state=seed, stratify=y)
    print(f"seed {seed}: fraud in test without stratify = {plain.sum()}, with stratify = {strat.sum()}")

seed 0: fraud in test without stratify = 3, with stratify = 2
seed 1: fraud in test without stratify = 2, with stratify = 2
seed 2: fraud in test without stratify = 2, with stratify = 2
seed 3: fraud in test without stratify = 4, with stratify = 2

Without stratification the fraud count jumps around, and so would any score you calculate on it. For classification, use stratify=y almost always. (*_, plain simply keeps the last returned array, y_test.)

Train, validation and test sets

While building a model you make many choices: which algorithm, which settings, which features. If you make those choices by checking the test score, you slowly fit the model to the test set and it stops being an honest exam. The fix is a third set:

You can make all three with two calls: first cut off the test set, then split the rest.

import numpy as np
from sklearn.model_selection import train_test_split

X = np.arange(1000).reshape(-1, 1)
y = np.arange(1000) % 2

X_rest, X_test, y_rest, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y)
X_train, X_val, y_train, y_val = train_test_split(
    X_rest, y_rest, test_size=0.25, random_state=42, stratify=y_rest)

print("train:", len(X_train), "validation:", len(X_val), "test:", len(X_test))

train: 600 validation: 200 test: 200

The second test_size is 0.25 because 25% of the remaining 800 rows is 200. With small datasets, a fixed validation set wastes data, so people use cross-validation instead, which rotates the validation part through the training data. You will learn it in Cross-Validation & Hyperparameter Tuning.

Time-ordered data

When rows are ordered in time, such as daily sales or stock prices, shuffling lets the model "see the future": it trains on next week's data and is tested on last week's. In real use you always predict forward, so test the same way: train on the past, test on the most recent rows, with shuffle=False.

import pandas as pd
from sklearn.model_selection import train_test_split

sales = pd.DataFrame({
    "day": pd.date_range("2026-10-01", periods=10, freq="D"),
    "boxes_sold": [42, 38, 45, 51, 60, 75, 49, 44, 47, 58],  # a Jaipur sweet shop
})
train, test = train_test_split(sales, test_size=0.3, shuffle=False)
print("Training days:", train["day"].min().date(), "to", train["day"].max().date())
print("Test days:    ", test["day"].min().date(), "to", test["day"].max().date())

Training days: 2026-10-01 to 2026-10-07
Test days:     2026-10-08 to 2026-10-10

For repeated time-based checks, scikit-learn provides TimeSeriesSplit, which always trains on earlier rows and tests on later ones.

Data leakage

Data leakage happens when information that would not be available at prediction time sneaks into training. It produces great test scores and disappointing real results. Common causes:

• Preparing data before splitting. Fitting a scaler or imputer on all rows lets test-set statistics shape training. Split first, then fit on the training set only.
• Duplicates across the split. If the same customer appears in both sets, the model is tested on rows it has already seen. Remove duplicates first, or keep each person's rows together with GroupShuffleSplit.
• Features from the future. Predicting loan default with a column like recovery_agent_assigned is cheating: that value only exists after the default happened.
• Shuffling time-ordered data, as shown above.

A useful habit: if a score looks too good to be true, look for leakage before celebrating.

Common mistakes

• Mixing up the return order. It is X_train, X_test, y_train, y_test. Writing X_train, y_train, X_test, y_test silently gives you wrong arrays.
• Forgetting random_state. Your scores change on every run and you cannot tell whether a change to the code helped.
• No stratify for classification. Small or imbalanced test sets can end up with too few examples of a class.
• Using the test set to choose settings. Use a validation set or cross-validation for decisions; open the test set once.

Key points

• Always measure a model on data it did not see during training; training scores can be perfect and meaningless.
• train_test_split returns train and test parts of each array you pass, in the order train, test.
• Set random_state for reproducible splits and stratify=y to keep class proportions.
• Use a validation set (or cross-validation) for model choices, and the test set only for the final score.
• For time-ordered data, use shuffle=False or TimeSeriesSplit so the model never trains on the future.
• Data leakage gives scores that are too good to be true; prepare data after splitting.

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

X, y = load_wine(return_X_y=True)
X_rest, X_test, y_rest, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y)
X_train, X_val, y_train, y_val = train_test_split(
    X_rest, y_rest, test_size=0.25, random_state=42, stratify=y_rest)

print("Sizes:", len(X_train), len(X_val), len(X_test))
print("Classes in full data:", np.bincount(y))
print("Classes in test set: ", np.bincount(y_test))

Sizes: 106 36 36
Classes in full data: [59 71 48]
Classes in test set:  [12 14 10]