Datasets & Splits

Train, validation, and test splits define what your model has seen during training versus what it must generalize to. Correct splitting is the foundation of honest evaluation.

Recipe

Quick-reference recipe card - copy-paste ready.

from sklearn.model_selection import train_test_split, StratifiedKFold, cross_val_score
 
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
scores = cross_val_score(model, X_train, y_train, cv=cv, scoring="f1_macro")

When to reach for this:

• Starting any supervised learning project before fitting preprocessors or models.
• Comparing models fairly with cross-validation on training data only.
• Handling imbalanced classes with stratified splits.
• Preventing leakage when multiple rows belong to the same user, session, or time window.

Working Example

"""datasets_splits.py - leak-free split workflow for tabular classification."""
from __future__ import annotations
 
import pandas as pd
from sklearn.datasets import load_breast_cancer
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import classification_report
from sklearn.model_selection import StratifiedKFold, cross_val_score, train_test_split
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler
 
# Load and inspect
data = load_breast_cancer(as_frame=True)
X: pd.DataFrame = data.data
y: pd.Series = data.target
 
# 1) Hold out final test set (never used during model selection)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
 
# 2) Build pipeline - fit only on training folds inside CV
pipe = Pipeline([
    ("scaler", StandardScaler()),
    ("clf", RandomForestClassifier(n_estimators=100, random_state=42)),
])
 
cv = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
cv_scores = cross_val_score(pipe, X_train, y_train, cv=cv, scoring="f1")
print(f"CV F1: {cv_scores.mean():.3f} +/- {cv_scores.std():.3f}")
 
# 3) Final fit on all training data, evaluate once on test
pipe.fit(X_train, y_train)
print(classification_report(y_test, pipe.predict(X_test)))

What this demonstrates:

• A locked test set evaluated only after model selection is complete.
• Cross-validation on training data for stable performance estimates.
• Preprocessing inside a pipeline so each CV fold fits the scaler independently.
• Stratified splitting to preserve class balance across folds.

Deep Dive

How It Works

• Train set fits model parameters and preprocessing statistics.
• Validation/CV folds tune hyperparameters and compare architectures without touching the test set.
• Test set provides a single unbiased estimate of generalization after all decisions are made.
• Stratification samples each class proportionally in every split, critical for imbalanced data.
• Group splits keep related rows (same patient, user, or day) in one partition only.

Split Strategies

Python Notes

from sklearn.model_selection import TimeSeriesSplit, GroupKFold
 
# Time series: always train on past, test on future
tscv = TimeSeriesSplit(n_splits=5)
 
# Groups: same group_id never in both train and test
gkf = GroupKFold(n_splits=5)
for train_idx, val_idx in gkf.split(X, y, groups=group_ids):
    ...

Gotchas

• Preprocessing before splitting - fitting a scaler on the full dataset leaks test statistics into training. Fix: split first, or wrap transforms in a Pipeline.
• Tuning on the test set - running grid search and picking the best model using test scores inflates metrics. Fix: use CV on training data; evaluate the winner once on test.
• Ignoring class imbalance - random splits can leave rare classes out of training folds. Fix: use stratify=y or resampling strategies.
• Shuffling time series - random splits mix future data into training for temporal problems. Fix: use TimeSeriesSplit or cutoff-based splits.
• Duplicate entities across splits - multiple rows per user in both train and test inflates scores. Fix: split by group_id, not by row.
• Small test sets - high variance in metrics with few test samples. Fix: increase test size or use repeated CV and report confidence intervals.

3730363230327c3234366f6967637c6f692e73656469756765646f63

Alternatives

FAQs

train_test_split(X, y, stratify=y)  # preserves class ratios in both partitions

train, test = train_test_split(df, test_size=0.2, stratify=df["label"])
X_train, y_train = train["features"], train["label"]

Related

• Feature Engineering - encoding after splitting
• Pipelines - leak-free preprocessing inside CV
• Model Evaluation - metrics on held-out data
• Imbalanced Data - stratified and resampled splits
• Hyperparameter Tuning - search inside CV folds

Stack versions: This page was written for Python 3.14.0 (stable 3.14, maintenance 3.13), FastAPI 0.115+, Django 5.2, Flask 3.1, Pydantic 2, PyTorch 2.6+, pandas 2.2+, Polars 1.x, ruff 0.9+, and uv 0.6+.