Stop 05 · Machine Learning FundamentalsTopics · 11/12

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.

The blob studies a big stack of practice sheets while a small sealed envelope of exam papers waits locked in a box.

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

r1
0
r2
1
r3
2
r4
3
r5
4
r6
5
r7
6
r8
7
r9
8
r10
9

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

python
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 number

1/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

~15 min

Quick check

Your model scores 99% on training data and 71% on test data. What is the most likely story?

Go deeper

  1. 1Practice `train_test_split` and understand the effect of the split ratio.
  2. 2Learn stratified splitting for classification tasks with imbalanced classes.
  3. 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.

Download
InsideEdge