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
| Strategy | Use When | sklearn API |
|---|---|---|
| Random split | IID rows, balanced classes | train_test_split |
| Stratified split | Classification with rare classes | stratify=y |
| Time-based split | Forecasting, temporal drift | TimeSeriesSplit |
| Group split | Multiple rows per entity | GroupKFold, GroupShuffleSplit |
| Nested CV | Unbiased hyperparameter tuning | GridSearchCV inside outer CV |
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=yor resampling strategies. - Shuffling time series - random splits mix future data into training for temporal problems. Fix: use
TimeSeriesSplitor 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.
Alternatives
| Alternative | Use When | Don't Use When |
|---|---|---|
| Single train/test split | Large datasets, quick iteration | Small data where split variance is high |
| K-fold cross-validation | Model comparison, limited data | Very large data where one split suffices |
| Holdout + validation set | Deep learning with long training runs | You need many hyperparameter trials (use CV) |
| Bootstrap / repeated CV | Unstable metrics on small data | Training cost per fold is prohibitive |
FAQs
How large should my test set be?
- 20% is a common default for datasets with thousands of rows.
- With fewer than 500 rows, prefer cross-validation over a large test holdout.
- Match test size to the minimum per-class count you need for reliable metrics.
When do I need a separate validation set if I already use cross-validation?
- CV on training data is enough for most classical ML hyperparameter tuning.
- A dedicated validation set helps when each training run is expensive (deep learning).
- Never use the test set as a validation set.
What does stratify=y do?
train_test_split(X, y, stratify=y) # preserves class ratios in both partitions- Each class appears in train and test at roughly the same proportion as the full dataset.
- Essential when any class is below 5% of rows.
Can I shuffle time series data before splitting?
- No - shuffling leaks future information into training.
- Use chronological splits or
TimeSeriesSplit. - Walk-forward validation is the production analog.
How do group splits work?
- Pass a
groupsarray with one ID per row (user_id, patient_id). GroupKFoldensures all rows for a group land in the same fold.- Prevents the model from memorizing entity-specific patterns that won't generalize.
What is nested cross-validation?
- Outer CV estimates generalization; inner CV tunes hyperparameters.
- Unbiased but computationally expensive.
- Use when publication-grade evaluation is required.
Should I set random_state?
- Yes -
random_state=42(any fixed int) makes splits reproducible. - Document the seed in experiment logs.
- Different seeds help sensitivity analysis on small datasets.
How do I split an already-loaded pandas DataFrame?
train, test = train_test_split(df, test_size=0.2, stratify=df["label"])
X_train, y_train = train["features"], train["label"]- Split the DataFrame first, then extract X and y.
- Keeps row indices aligned between features and labels.
What is data leakage in the context of splits?
- Any information from test data influencing training decisions.
- Common sources: global scaling, feature selection on full data, duplicate entities.
- Pipelines and proper split order prevent most leakage.
When is cross_val_score not enough?
- When you need per-fold predictions for calibration or error analysis.
- Use
cross_val_predictinstead, still only on training data. - For final metrics, fit on full training set and predict test once.
How do I handle multiple target columns?
- Multi-output problems need
MultiOutputClassifieror separate models per target. - Stratification works on a single label column - bin combined labels or stratify the primary target.
Can I reuse the test set across projects?
- Yes, if the test set is truly held out and never used for any tuning.
- Document that the test set is locked and shared across experiments.
- Consider a fresh test set for major data distribution shifts.
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+.