Train/Test and Train/Validation/Test Split

scikit-learn's train_test_split is the standard tool for splitting data.

Basic train/test split

import pandas as pd
from sklearn.model_selection import train_test_split

df = pd.read_csv("Titanic-Dataset.csv")
X = df[["Pclass", "Age", "SibSp", "Parch", "Fare"]]
y = df["Survived"]

X_train, X_test, y_train, y_test = train_test_split(
    X, y,
    test_size=0.2,       # 20% for testing
    random_state=42,     # reproducible
    shuffle=True,
    stratify=y,          # keep class ratio
)

print(X_train.shape, X_test.shape)

Key parameters

Parameter	Meaning
test_size / train_size	Fraction (0.2 = 20%) or an absolute number of rows
random_state	Fixes the random seed so the split is reproducible
shuffle	Mix the rows before splitting (default True; set False for time series)
stratify	Keeps the class proportions the same in every split

Why stratify=y matters

Without stratification, a random split of imbalanced data may leave the test set with very few (or zero) minority examples.

print(y.value_counts(normalize=True).round(3))
print(y_train.value_counts(normalize=True).round(3))
print(y_test.value_counts(normalize=True).round(3))   # nearly identical thanks to stratify

Train / validation / test split

train_test_split splits into two parts, so call it twice.

# 1st split: hold out the test set (15%)
X_temp, X_test, y_temp, y_test = train_test_split(
    X, y, test_size=0.15, stratify=y, random_state=42
)

# 2nd split: validation is ~15% of the whole data
X_train, X_val, y_train, y_val = train_test_split(
    X_temp, y_temp, test_size=0.1765, stratify=y_temp, random_state=42
)

The test_size=0.1765 in the second call is 15 / 85, which gives 15% of the original data.

Split sizes at a glance

print("Train:", len(X_train), "| Validation:", len(X_val), "| Test:", len(X_test))

Which split is used for what?

1. Train — model.fit(X_train, y_train)
2. Validation — compare models and tune hyperparameters
3. Test — evaluate once at the very end