Test a model on the rows it trained on and it will flatter itself every time. A train test split in Python fixes that: hide a slice of the data, train on the rest, and let the hidden slice deliver the only score you can trust. This guide covers the basic split, stratification, K-Fold cross-validation, and the three-way validation setup, all with runnable scikit-learn code.
“Testing on your training data is like grading your own exam. You will always think you did great.”
Sebastian Raschka, Python Machine Learning
Last Updated: July 2026 | Tested on: Python 3.14.6, scikit-learn 1.9.0 | Difficulty: Intermediate | Reading Time: 16 minutes
Think about how a good teacher writes an exam. The practice questions you study from are not the questions on the final. If they were, everyone would score 100% and the marks would tell you nothing. That gap between what you studied and what you are tested on is the whole point of a train test split in Python. You teach the model on one slice of the data, then judge it on a slice it has never seen. The slice you hide away is the test set, and it is the only honest measure of how the model will do on tomorrow’s real data.
Here is the catch. One single split is shaky. Depending on which rows happen to fall into the test set, your accuracy can jump around by several percent. It is a bit like judging a restaurant from one meal. You might have caught the chef on a great night, or a terrible one. To get a fair verdict you need to eat there a few times. Cross-validation does exactly that for models: it splits the data several different ways, scores each one, and averages the results so you get a number you can trust.
A junior data scientist named Pravin lived this. He trained a model, tested it, and proudly reported 92% accuracy. His manager asked him to retrain with a different random seed, and the new split gave 84%. So which number is real? Neither. They are both accidents of how the rows landed. Cross-validation would have handed Pravin something like “88% give or take 3%”, and that range is the honest answer he should have reported in the first place.
Table of Contents
Prerequisites
- feature selection tutorial
- machine learning introduction (the overall ML workflow)
The Basic Train/Test Split
Tip: click the diagram to open it full screen, then use the zoom and pan controls for a closer look.
The diagram above shows the most common way to carve up a dataset: 70% to train on (the model learns from this), 15% for validation (where you compare settings and pick the best one), and 15% for the final test (the one honest score you report). Think of the validation set as your practice exams and the test set as the real final. You sit as many practice exams as you like, but you only get one shot at the final, and you do not get to peek at it first.
Splitting the data this way is what stops you from quietly tuning your model until it happens to look good on the very data you are supposed to be judging it with.
📄 basic_split.py: train/test split fundamentals
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
X, y = make_classification(n_samples=1000, n_features=10, random_state=42)
# Basic split: 80% train, 20% test
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=42
)
print(f"Total samples: {len(X)}")
print(f"Train: {len(X_train)} ({len(X_train)/len(X):.0%})")
print(f"Test: {len(X_test)} ({len(X_test)/len(X):.0%})")
# Train and evaluate
model = LogisticRegression(random_state=42)
model.fit(X_train, y_train)
print(f"\nTrain accuracy: {model.score(X_train, y_train):.3f}")
print(f"Test accuracy: {model.score(X_test, y_test):.3f}")
# Different random seeds give different results!
for seed in [0, 1, 2, 3, 42]:
Xtr, Xte, ytr, yte = train_test_split(X, y, test_size=0.2, random_state=seed)
m = LogisticRegression(random_state=42).fit(Xtr, ytr)
print(f"Seed {seed:2d}: test accuracy = {m.score(Xte, yte):.3f}")
▶ Output
Total samples: 1000 Train: 800 (80%) Test: 200 (20%) Train accuracy: 0.870 Test accuracy: 0.830 Seed 0: test accuracy = 0.855 Seed 1: test accuracy = 0.865 Seed 2: test accuracy = 0.875 Seed 3: test accuracy = 0.865 Seed 42: test accuracy = 0.830
What happened here: Same model, same data, yet the test accuracy wanders from 0.830 all the way up to 0.875 just because we changed the random seed. That is a swing of about 4.5 points, and nothing about the model changed. The only thing that moved was which 200 rows happened to land in the test set. That single fact is why you should never trust one train/test split on its own, and it is exactly the problem cross-validation was built to fix. (Your own numbers may land a little differently across scikit-learn versions, but the wobble is always there.)
Stratified Splitting: Keeping Class Balance
A plain random split treats every row the same, which is fine until your classes are lopsided. Imagine a fraud dataset where only 5 in every 100 transactions are fraud. If the random shuffle is unlucky, your test set might end up with barely any fraud cases in it, and then your “fraud detector” was never really tested on fraud at all. Stratified splitting fixes this by keeping the same class proportions in every slice. It is like splitting a class of 30 students into two study groups but making sure each group keeps the same mix of beginners and experts, instead of accidentally dumping all the experts into one group.
📄 stratified_split.py: why stratification matters for imbalanced data
import numpy as np
from sklearn.model_selection import train_test_split
rng = np.random.default_rng(42)
# Imbalanced dataset: 95% class 0, 5% class 1 (like fraud detection)
y_imbalanced = np.array([0] * 950 + [1] * 50)
X_dummy = rng.normal(size=(1000, 5))
# Without stratification, class 1 might be underrepresented
_, _, _, y_test_bad = train_test_split(X_dummy, y_imbalanced, test_size=0.2, random_state=0)
print(f"Without stratification:")
print(f" Test set class 1: {(y_test_bad == 1).sum()}/{len(y_test_bad)} = {(y_test_bad == 1).mean():.1%}")
# With stratification, proportions are preserved
_, _, _, y_test_good = train_test_split(X_dummy, y_imbalanced, test_size=0.2,
random_state=0, stratify=y_imbalanced)
print(f"With stratification:")
print(f" Test set class 1: {(y_test_good == 1).sum()}/{len(y_test_good)} = {(y_test_good == 1).mean():.1%}")
print(f" Original class 1: {(y_imbalanced == 1).sum()}/{len(y_imbalanced)} = {(y_imbalanced == 1).mean():.1%}")
▶ Output
Without stratification: Test set class 1: 12/200 = 6.0% With stratification: Test set class 1: 10/200 = 5.0% Original class 1: 50/1000 = 5.0%
What happened here: The plain random split landed 12 fraud cases in the test set, which is 6%, when the true rate is only 5%. It overshot this time. With a different seed it could just as easily undershoot and leave you with 3% or 4%. Either way the test set no longer reflects reality. Add stratify=y and the split nails exactly 5% in every slice, matching the original distribution. The rule of thumb is simple: any time your classes are imbalanced, pass stratify=y so the test set stays a fair sample of the real world.
K-Fold Cross-Validation
Here is the fix for that “one shaky split” problem. Think of how a cook tastes a big pot of dal. Dipping the spoon in just once, right at the top, tells you almost nothing. A careful cook stirs and tastes from a few different spots, then trusts the overall impression. K-Fold does the same with your data. Instead of betting everything on a single split, K-Fold chops the data into K equal slices, called folds.
It then trains K separate models. Each one holds out a different fold as the validation set and trains on the other K minus 1 folds. Rotate through all the folds and every single row gets a turn at being validated exactly once, with no row ever wasted. Your final score is just the average of the K fold scores, which is far steadier than any one number on its own.
Tip: click the diagram to open it full screen, then use the zoom and pan controls for a closer look.
The diagram illustrates K-Fold cross-validation with 5 folds: the data is split into 5 equal parts, and in each round a different fold serves as the validation set while the remaining 4 folds are used for training. This rotation ensures every data point is used for both training and validation exactly once. The final score is the average across all 5 folds, giving a more robust estimate of model performance than a single train/test split, especially when data is limited.
📄 cross_validation.py: K-Fold and stratified K-Fold
import numpy as np
from sklearn.model_selection import cross_val_score, KFold, StratifiedKFold
from sklearn.datasets import make_classification
from sklearn.linear_model import LogisticRegression
X, y = make_classification(n_samples=1000, n_features=10, random_state=42)
model = LogisticRegression(random_state=42)
# Simple cross_val_score, the most common usage
scores = cross_val_score(model, X, y, cv=5, scoring="accuracy")
print(f"5-Fold CV scores: {scores.round(3)}")
print(f"Mean: {scores.mean():.3f} ± {scores.std():.3f}")
# Stratified K-Fold (default for classification in cross_val_score)
skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
strat_scores = cross_val_score(model, X, y, cv=skf, scoring="accuracy")
print(f"\nStratified 5-Fold: {strat_scores.round(3)}")
print(f"Mean: {strat_scores.mean():.3f} ± {strat_scores.std():.3f}")
# Compare different K values
for k in [3, 5, 10]:
cv_scores = cross_val_score(model, X, y, cv=k)
print(f"\n{k:2d}-Fold: mean={cv_scores.mean():.3f} ± {cv_scores.std():.3f}")
▶ Output
5-Fold CV scores: [0.87 0.855 0.85 0.83 0.875] Mean: 0.856 ± 0.016 Stratified 5-Fold: [0.83 0.84 0.875 0.89 0.845] Mean: 0.856 ± 0.023 3-Fold: mean=0.854 ± 0.015 5-Fold: mean=0.856 ± 0.016 10-Fold: mean=0.857 ± 0.030
What happened here: Instead of one nervous single number, you now get “0.856 give or take 0.016”, a mean with a spread around it. That spread is the honest part. It tells you how much the score bounces depending on which rows the model trained on. One detail worth knowing: when you pass cv=5 to cross_val_score on a classifier, scikit-learn already uses a stratified split under the hood, just without shuffling first. The second block builds its own StratifiedKFold(shuffle=True, random_state=42), and shuffling reshuffles which rows sit together, so the two runs land on different fold scores (here the shuffled one even spreads a bit wider).
Both are valid. The takeaway is that K=5 is the sensible default. K=10 trains on slightly more data per fold but costs more compute and, as the numbers show, can jiggle more from fold to fold. Your exact decimals will shift across scikit-learn versions, so trust the pattern, not the third digit.
Train, Validation, Test: The Three-Way Split
The moment you start tuning hyperparameters, two sets are no longer enough. You need three. The training set teaches the model. The validation set is where you try out different settings and pick the winner. The test set stays sealed until the very end and gives you the one final score you report. Here is why the third set matters: every time you tweak a setting because it scored well on a particular set, you are quietly learning the quirks of that set. Do that on your test set and your reported accuracy turns into wishful thinking. It is the difference between a mock exam you can retake and the real one you sit only once.
📄 three_way_split.py: a clean hyperparameter tuning workflow
import numpy as np
from sklearn.model_selection import train_test_split, cross_val_score
from sklearn.datasets import make_classification
from sklearn.ensemble import RandomForestClassifier
X, y = make_classification(n_samples=1000, n_features=10, random_state=42)
# Step 1: Hold out test set (untouched until the very end)
X_dev, X_test, y_dev, y_test = train_test_split(
X, y, test_size=0.15, random_state=42, stratify=y
)
print(f"Development set: {len(X_dev)} | Test set: {len(X_test)} (sealed!)")
# Step 2: Use cross-validation on dev set to choose hyperparameters
best_score, best_n = 0, 0
for n_trees in [50, 100, 200, 500]:
scores = cross_val_score(
RandomForestClassifier(n_estimators=n_trees, random_state=42),
X_dev, y_dev, cv=5
)
mean = scores.mean()
print(f" n_trees={n_trees:3d}: CV score = {mean:.3f} ± {scores.std():.3f}")
if mean > best_score:
best_score, best_n = mean, n_trees
print(f"\nBest hyperparameter: n_trees={best_n}")
# Step 3: Train final model on ALL dev data with best hyperparameters
final_model = RandomForestClassifier(n_estimators=best_n, random_state=42)
final_model.fit(X_dev, y_dev)
# Step 4: Evaluate ONCE on test set, this is the true performance
test_score = final_model.score(X_test, y_test)
print(f"Final test score: {test_score:.3f}")
print("(This number is reported. The test set was never used for tuning.)")
▶ Output
Development set: 850 | Test set: 150 (sealed!) n_trees= 50: CV score = 0.908 ± 0.015 n_trees=100: CV score = 0.911 ± 0.012 n_trees=200: CV score = 0.912 ± 0.013 n_trees=500: CV score = 0.911 ± 0.010 Best hyperparameter: n_trees=200 Final test score: 0.907 (This number is reported. The test set was never used for tuning.)
What happened here: Notice the order of operations. We sealed the 150-row test set first, before touching anything. Then we did all the comparing of 50, 100, 200, and 500 trees using cross-validation on the development set only. The test set never saw any of that. Only after picking 200 trees did we open the envelope once and read 0.907. That single number is the one you put in your report, and you can trust it precisely because the test set had no say in choosing the model.
Also worth noticing: the cross-validation (CV) scores for 100, 200, and 500 trees are almost identical (0.911 to 0.912), which is a gentle hint that piling on more trees past a point buys you basically nothing.
Common Mistakes
❌ Mistake 1: Tuning hyperparameters on the test set
# BAD: Testing different hyperparameters on test set
# for C in [0.01, 0.1, 1, 10]:
# model = LogisticRegression(C=C).fit(X_train, y_train)
# print(model.score(X_test, y_test)) # Leaking test info!
# Each time you evaluate on the test set, you learn something about it.
# Choosing the C that scores best on the test set = training on test data.
# GOOD: Use cross-validation on training data to choose C.
# Evaluate on test set ONCE, with the chosen C.
print("The test set is a sealed envelope. Open it once, at the end.")
❌ Mistake 2: Shuffling time-series data
# BAD: Random split on time-series lets future data leak into training
# train_test_split(stock_data, shuffle=True)
# GOOD: Time-based split, train on the past, test on the future
# cutoff = int(len(data) * 0.8)
# train, test = data[:cutoff], data[cutoff:]
# Or use TimeSeriesSplit from sklearn
from sklearn.model_selection import TimeSeriesSplit
tscv = TimeSeriesSplit(n_splits=5)
print("TimeSeriesSplit: always trains on past, tests on future")
for i, (train_idx, test_idx) in enumerate(tscv.split(range(100))):
print(f" Fold {i}: train {train_idx[0]}-{train_idx[-1]}, test {test_idx[0]}-{test_idx[-1]}")
▶ Output
TimeSeriesSplit: always trains on past, tests on future Fold 0: train 0-19, test 20-35 Fold 1: train 0-35, test 36-51 Fold 2: train 0-51, test 52-67 Fold 3: train 0-67, test 68-83 Fold 4: train 0-83, test 84-99
What happened here: Look at the fold boundaries. The training window only ever grows forward in time and the test fold always sits right after it, never before. That is the whole point. If you shuffled stock prices or sales data the normal way, the model could end up training on next March while being tested on last January, which is basically letting it peek at the future. No real system gets to do that. TimeSeriesSplit keeps the arrow of time pointing one way, the same way you would never grade a student on questions they had already seen the answers to.
Practice Exercises
- Exercise 1: Make an 80/20 split with
train_test_split, then print the size of each set and confirm the split really is 80/20. - Exercise 2: Build a dataset that is 90% one class and 10% the other, split it with and without
stratify=y, and compare how close each test set comes to the true 10%. - Exercise 3: Run 5-fold cross-validation on the same model three ways: plain K-Fold, stratified K-Fold, and
TimeSeriesSplit. Print the mean and standard deviation each time and explain why the time-series numbers look different.
Conclusion
You now have the honest toolkit for judging a model. You learned why a single train test split wobbles from seed to seed, how stratify=y keeps rare classes fairly represented, how K-Fold cross-validation turns one nervous number into a mean with a spread you can trust, and why a sealed three-way split (train, validation, test) is the only way to report accuracy you can actually defend. You also saw the traps: tuning on the test set and shuffling time-series data, both of which quietly leak the answers.
The one habit to carry forward: pick your model with cross-validation, and open the test set exactly once, at the very end. Next up, we clean the fuel that feeds every model in ML: Data Preprocessing, Your Model Is Only as Good as Your Data, where honest evaluation meets the raw, messy data it has to score. For the full roadmap from beginner to professional, visit the Python + AI/ML tutorial series home.
Frequently Asked Questions
What is the best train/test split ratio?
80/20 or 70/30 for datasets with 1,000-10,000 samples. For very large datasets (100K+), even 90/10 or 95/5 works because the test set is still large enough. For very small datasets (<500), skip the fixed split entirely and use Leave-One-Out cross-validation.
Should I use 5-fold or 10-fold cross-validation?
5-fold is the standard. It balances computational cost (trains 5 models) with statistical reliability. 10-fold uses more training data per fold, giving slightly lower bias but higher variance and 2x the computation. For large datasets, 5-fold is sufficient.
What does the standard deviation in cross-validation mean?
It measures how much your model’s performance varies across different data subsets. A high standard deviation (say 0.860 ± 0.05) means the model is sensitive to which data it sees, so you might need more data or a more stable algorithm. A low standard deviation (0.860 ± 0.01) means robust, reliable performance.
Can I use the test set more than once?
Ideally no. Every time you evaluate on the test set and then change your model, you are implicitly training on it. In practice, most people peek at the test set a few times. The key is: do not choose between models based on test set performance. Use cross-validation for model selection, test set for final reporting.
How do I do a train test split in Python?
Use train_test_split from scikit-learn: X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42). Set test_size to the fraction you want to hold out, fix random_state so the split is reproducible, and add stratify=y when your classes are imbalanced. That one call is the standard train test split in Python, and nearly every model starts with it.
Interview Questions on Train Test Split and Cross-Validation
How interviewers actually probe this topic: real scenarios, with answers you can say out loud.
Q: What is the purpose of the random_state argument in train_test_split, and when does it matter?
It fixes the shuffle so the same rows land in the same set every run, which makes your results reproducible. That matters when you are debugging, comparing two models fairly, or sharing code so a teammate gets identical splits. It does not make one split “correct”; different seeds still give different scores, which is exactly why cross-validation exists.
Q: Why is testing a model on its training data misleading?
The model has already seen those exact rows, so it can score high just by memorizing them rather than learning a general pattern. That inflated number tells you nothing about tomorrow’s unseen data. A held-out test set is the only honest measure of how the model will actually perform in production.
Q: When should you pass stratify=y, and what happens if you forget it?
Use it whenever classes are imbalanced, so each split keeps the same class proportions as the full dataset. If you forget it on, say, a fraud dataset with 5% fraud, an unlucky shuffle can leave the test set with almost no fraud cases, so your “fraud detector” is barely tested on fraud at all. Stratification keeps every slice a fair sample of the real distribution.
Q: What is the difference between the validation set and the test set?
The validation set is where you compare settings and pick the best model, so you use it many times. The test set stays sealed until the very end and gives you one final score you report. Mixing them up (tuning against the test set) leaks information and turns your reported accuracy into wishful thinking.
Q: A teammate named Anvi reports 95% test accuracy on a fraud model, but it flags almost nothing in production. What do you check first?
First check whether the split was stratified: with heavy imbalance an unstratified test set can be nearly all the majority class, so 95% just means “predict not-fraud every time.” Then look for data leakage, such as a feature that encodes the label or preprocessing (scaling, resampling) fit on the full data before splitting. Finally, confirm she is looking at the right metric; accuracy is a poor choice for imbalanced data, where precision, recall, or Receiver Operating Characteristic (ROC)-Area Under the Curve (AUC) tell the real story.
Q: You are building a model to forecast daily sales and your cross-validation scores look great, but live predictions are poor. What is the likely cause?
You almost certainly shuffled time-series data, which lets the model train on future dates and test on past ones, effectively peeking at the answer. Switch to a time-aware split like TimeSeriesSplit, which always trains on the past and tests on the future. Also make sure any feature engineering (rolling averages, lag features) never pulls in values from after the prediction date.
Q: Why might 10-fold cross-validation show a higher standard deviation than 5-fold on the same data?
With 10 folds each validation fold is smaller, so a few tricky rows swing its score more, and averaging over smaller test sets makes the fold-to-fold spread wider. You train on slightly more data per fold, which can lower bias, but the smaller validation slices add variance. That trade-off is why 5-fold is the common default: steadier estimates at lower compute cost.
Series: Python + AI/ML Cookbook, Part 5: Machine Learning
Go deeper: the official Python documentation covers every edge case of this topic.
Related Posts
Previous: ML: Setting Up ML Environment with Jupyter and scikit-learn
Next: ML: Data Preprocessing, Your Model Is Only as Good as Your Data
Series Home: Python + AI/ML Tutorial Series

No comment