Train/Test Split
Part of the AI Engineer Roadmap
In short
A model that scores 100% on data it has already seen tells you nothing. Holding back a test set is how you find out if it actually learnt anything.
Holding out a portion of your data to evaluate a model on examples it never trained on — the basic safeguard against fooling yourself with an overfit model.

Think of it like this
The exam with leaked questions
If the exam paper has the exact questions from your practice sheet, a student who mugged up answers will top the class. A fair exam uses fresh questions. Your test set is that fresh paper, kept locked away until the very end.
Watch it happen
Splitting 10 rows 80/20
1/4Ten labelled rows, in the order they were collected. Maybe the last ones are all from December.
The big ideas
What you need to know
Test data stays unseen
Never fit, tune or peek at the test set. Once you use it to make choices, its score stops being honest.
random_state for repeatability
The split is random. Setting random_state=42 (any number works) gives the same split every time you run.
stratify keeps the balance
If only 10% of rows are fraud, stratify=y makes sure both sets keep roughly that 10%.
Time data splits by date
For sales or stock data, train on the past and test on the future. Shuffling would let the model peek ahead.
See it in code
Walk through the code
1from sklearn.datasets import load_wine2from sklearn.model_selection import train_test_split3from sklearn.tree import DecisionTreeClassifier4 5X, y = load_wine(return_X_y=True)6X_train, X_test, y_train, y_test = train_test_split(7 X, y, test_size=0.2, stratify=y, random_state=428)9 10tree = DecisionTreeClassifier(random_state=0).fit(X_train, y_train)11print("train:", tree.score(X_train, y_train)) # 1.0, it memorised12print("test: ", tree.score(X_test, y_test)) # lower, the honest number1/320% of rows go to the test set, with the same class mix as the full data.
A deep decision tree aces the rows it trained on. The test score is the one you report.
Save yourself the pain
Common mistakes
Oops
Trying ten models and picking the one with the best test score.
Do this instead
Compare models with cross-validation or a separate validation set. Touch the test set once, at the end.
Oops
Cleaning or scaling the full dataset before splitting.
Do this instead
Split first. Anything you learn from data, like averages for filling blanks, should come from training rows only.
Your mission
Catch a memoriser
Quick check
Your model scores 99% on training data and 71% on test data. What is the most likely story?
Go deeper
- 1Practice `train_test_split` and understand the effect of the split ratio.
- 2Learn stratified splitting for classification tasks with imbalanced classes.
- 3Understand why you must never evaluate on data the model has seen during training.
More topics at this stop
Stuck on this topic? Ask an Insider
Get 1:1 guidance from people who've walked this exact path — on the InsideEdge app.