Train Test Split in Machine Learning
A train test split divides a dataset into two parts: a training set used to fit the model and a test set used to evaluate how well the model performs on unseen data. In Python, this is commonly done with sklearn.model_selection.train_test_split().
Quick sklearn train_test_split example — 80/20 split
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(
X,
y,
test_size=0.20,
random_state=42,
stratify=y
)Here, 80% of the rows go to training and 20% go to testing.random_state=42 makes the split reproducible, while stratify=y helps preserve class proportions for classification.
Train Test Split in Simple Words
1. Start with data
100 available examples
2. Train
Use, for example, 80 examples
3. Test
Evaluate on the remaining 20
Main idea: the test set acts like a small simulation of future unseen data.
Primary Purpose
Estimate how well a trained model generalizes beyond the examples used to fit it.
Why Train Test Split Is Necessary
Suppose a student practices using one set of questions and is then evaluated using different questions. A high score on the unseen questions gives stronger evidence that the student learned the concept rather than only remembering the practice answers.
- Practice questions → Training Set
- Final unseen questions → Test Set
Machine Learning follows the same idea. Evaluating on the same examples used for fitting usually gives an overly optimistic picture of performance. A held-out test set helps measure generalization. It can reveal overfitting, but the act of splitting data by itself does not guarantee that overfitting is prevented.
A Tiny 10-Student Example
Training set — 8 students
S1 S2 S3 S4 S5 S6 S7 S8
The model is allowed to learn from these rows.
Test set — 2 students
S9 S10
These rows are held back for evaluation.
Common Split Ratios — Examples, Not Rules
| Example | Training | Testing |
|---|---|---|
| 70 / 30 | 70% | 30% |
| 80 / 20 | 80% | 20% |
| 90 / 10 | 90% | 10% |
There is no universal best ratio. The choice depends on dataset size, how much data the model needs to learn, and how precise you need the evaluation to be. An 80/20 split is a common teaching choice, not a law.
Train, Validation & Test Sets
Training Set
Used to fit model parameters from data.
Validation Set / CV
Used for model comparison, hyperparameter choices and other development decisions.
Test Set
Kept aside for the final estimate of performance on unseen data.
Why Validation Is Different from Testing
If you repeatedly check the test score while choosing models or hyperparameters, information from the test set begins influencing your decisions. The test score can then become optimistic. Use a separate validation set or cross-validation for model selection, and reserve the test set for final evaluation.
A Safer ML Workflow
Simple Split Mathematics
Suppose the dataset contains N = 1000 samples and test_size=0.2.
Test samples = 1000 × 0.2 = 200
Training samples = 1000 − 200 = 800
Random, Stratified, Grouped & Time-Aware Splits
Random Splitting
For independent rows, shuffling before a split can prevent accidental ordering from putting systematically different observations into train and test sets. Scikit-learn's train_test_split() shuffles by default.
random_state
Setting random_state=42 makes the pseudo-random split reproducible when the other inputs are unchanged. It does not make the model more accurate; it makes the experiment repeatable.
Stratified Splitting for Classification
When class proportions matter, stratify=y can preserve approximately the same class proportions in train and test sets.
Stratification is useful for many classification problems, especially when a class is uncommon. It is not a universal instruction to use on every dataset.
Grouped Data
If several rows belong to the same person, patient, machine, household or customer, placing the same entity in both train and test sets can leak entity-specific information. In such cases, split by group rather than by individual rows.
Time-Ordered Data
For forecasting, randomly shuffling past and future observations can create an unrealistic evaluation. A time-aware split generally trains on earlier observations and evaluates on later observations.
Data Leakage: Split Before Learning from Data
Data leakage happens when information that should be unavailable during training influences the fitted model or model-selection process. A common beginner mistake is fitting preprocessing on the complete dataset before the split.
Wrong
Better
The same principle applies to imputers, feature selection, PCA and many other transformations. In real projects, a Scikit-learn Pipeline is often a safer way to keep preprocessing and model fitting inside the correct training workflow.
sklearn train_test_split in Python
Scikit-learn provides train_test_split() in sklearn.model_selection. The main arguments beginners usually need are test_size, random_state, and, for many classification problems, stratify.
Python Code: Split, Train and Evaluate
import numpy as np
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score
# Step 1 — Load a small classification dataset
X, y = load_iris(return_X_y=True)
# Step 2 — Hold out 20% for testing
X_train, X_test, y_train, y_test = train_test_split(
X,
y,
test_size=0.20,
random_state=42,
shuffle=True,
stratify=y
)
print("Train shape:", X_train.shape)
print("Test shape:", X_test.shape)
print("Train class counts:", np.bincount(y_train))
print("Test class counts:", np.bincount(y_test))
# Step 3 — Train only on the training set
model = LogisticRegression(max_iter=1000)
model.fit(X_train, y_train)
# Step 4 — Evaluate on the held-out test set
y_pred = model.predict(X_test)
accuracy = accuracy_score(y_test, y_pred)
print("Test accuracy:", round(accuracy, 3))The equal class counts are a direct consequence of using stratify=y on the balanced Iris dataset with this split size. The accuracy is specific to this dataset and split; it is not a general promise for Logistic Regression.
Common Mistakes Beginners Make
- Training and testing on the same rows: gives an unreliable estimate of future performance.
- Fitting preprocessing before the split: leaks information from held-out data.
- Using the test set for hyperparameter tuning: gradually turns the test set into part of the development process.
- Always shuffling: can be wrong for time-ordered data.
- Always stratifying: is not appropriate for every problem; it is mainly a classification tool.
- Ignoring groups: can place the same person or entity in both train and test sets.
- Using a tiny test set: can make the estimate highly unstable.
Train/Test Split vs Cross-Validation
| Feature | Single Train/Test Split | Cross-Validation |
|---|---|---|
| Number of evaluations | One split | Multiple folds/splits |
| Computation | Usually lower | Usually higher |
| Sensitivity to one split | Higher | Can reduce dependence on one split |
| Typical use | Quick holdout / final test | Model comparison and tuning |
Advantages
- Simple to understand and implement.
- Usually cheaper than repeated cross-validation.
- Useful when the dataset is large enough to reserve a meaningful holdout set.
- Provides a clear untouched set for final evaluation.
Limitations
- Performance can depend noticeably on which rows land in the split.
- Holding out data can be costly when the dataset is very small.
- A naive random split can be invalid for time, groups or other dependencies.
Quick Recap
Why should the test set stay out of model tuning?
Because repeated decisions based on test performance leak information from the test set into model development, making the final test score less trustworthy.
What does random_state=42 do?
It makes the pseudo-random split reproducible. The number 42 itself has no special statistical advantage.
Should time-series data be randomly shuffled?
Usually no. A forecasting evaluation should respect time order so the model learns from the past and is evaluated on later observations.
Continue Learning
A single holdout split is only the beginning of model evaluation.
Final Summary
Train/Test Split divides available data so that model fitting and final evaluation are separated. The training set teaches the model; validation or cross-validation supports development decisions; the test set estimates performance on held-out data.
In Scikit-learn, train_test_split() provides a convenient random holdout split with options such as test_size, random_state, shuffle and stratify. The correct splitting strategy, however, depends on the structure of the data: ordinary random splitting is not automatically appropriate for time-series or grouped observations.
Most Important Insight to Remember
Keep the final test set outside model fitting and tuning. Split first, learn preprocessing from training data, use validation or cross-validation for choices, and use the test set only when you are ready for a final evaluation.