Machine learningChapter 104 of 114
Train and Test Split
Hold data back, or you cannot tell whether the model learned anything.
Why hold data back
A model that has seen an example can simply remember it. Scoring on the data it trained on tells you about its memory, not its ability:
from sklearn.datasets import load_breast_cancer
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier
X, y = load_breast_cancer(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=0, stratify=y
)
model = DecisionTreeClassifier(random_state=0).fit(X_train, y_train)
print("on data it trained on:", round(model.score(X_train, y_train), 2))
print("on data it never saw: ", round(model.score(X_test, y_test), 2))Output
on data it trained on: 1.0 on data it never saw: 0.94
A perfect score on the training set is normal for a decision tree and means nothing at all — it grew branches until every training example was placed correctly. Only the second number is a claim about the future.
The call
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
X, y = load_iris(return_X_y=True)
X_train, X_test, y_train, y_test = train_test_split(
X, y, test_size=0.2, random_state=0
)
print(X_train.shape, X_test.shape)
print(y_train.shape, y_test.shape)Output
(120, 4) (30, 4) (120,) (30,)
The return order catches everyone once: X_train, X_test, y_train, y_test. All four, X before y, train before test.
test_size
A fraction, or a count:
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
X, y = load_iris(return_X_y=True)
for size in [0.2, 0.3, 50]:
_, X_test, _, _ = train_test_split(X, y, test_size=size, random_state=0)
print(f"test_size={size!r:5} -> {X_test.shape[0]} test rows")Output
test_size=0.2 -> 30 test rows test_size=0.3 -> 45 test rows test_size=50 -> 50 test rows
Twenty to thirty percent is the usual choice. Too small a test set and the score is noise; too large and the model has little to learn from.
random_state
The split is random. Without a seed you get a different one every run, and a different score:
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
X, y = load_iris(return_X_y=True)
a = train_test_split(X, y, test_size=0.2, random_state=0)[1]
b = train_test_split(X, y, test_size=0.2, random_state=0)[1]
c = train_test_split(X, y, test_size=0.2, random_state=1)[1]
print((a == b).all())
print((a == c).all())Output
True False
Set random_state in anything you want to reproduce or compare. Leave it out only when you are deliberately measuring how much the split matters.
stratify
By default the split is uniformly random, which can badly skew a small or imbalanced dataset. stratify=y keeps the class proportions:
import numpy as np
from sklearn.model_selection import train_test_split
y = np.array([0] * 90 + [1] * 10)
X = np.arange(100).reshape(-1, 1)
_, _, _, plain = train_test_split(X, y, test_size=0.2, random_state=3)
_, _, _, kept = train_test_split(X, y, test_size=0.2, random_state=3, stratify=y)
print("without stratify, rare class in test:", int((plain == 1).sum()))
print("with stratify, rare class in test:", int((kept == 1).sum()))Output
without stratify, rare class in test: 4 with stratify, rare class in test: 2
Ten percent of the data is class 1, so a fair test set of 20 holds 2. This split happened to put 4 there — double the share — leaving the training set short of exactly the class it most needs to learn. Another seed would take too few instead. Use stratify=y for classification as a matter of habit.
The test set is not for tuning
Once you start choosing settings by test score, the test set has quietly become part of training and its number is optimistic. The honest arrangement is three groups:
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
X, y = load_iris(return_X_y=True)
X_rest, X_test, y_rest, y_test = train_test_split(
X, y, test_size=0.2, random_state=0, stratify=y
)
X_train, X_val, y_train, y_val = train_test_split(
X_rest, y_rest, test_size=0.25, random_state=0, stratify=y_rest
)
print("train", X_train.shape[0], "validation", X_val.shape[0], "test", X_test.shape[0])Output
train 90 validation 30 test 30
Tune against the validation set. Touch the test set once, at the end, to report a number. The Overfitting chapter shows cross-validation, which does this more efficiently when data is scarce.
Test yourself
3 questionsWhy is a score on the training data misleading?
Show the answer
The model may simply have memorised those examples — A perfect training score is normal for a decision tree and tells you nothing about new data.
What does stratify=y do?
Show the answer
Keeps the class proportions the same in both halves — Reproducibility is random_state. Use stratify for classification as a matter of habit.
What is wrong with trying several random_state values and keeping the best score?
Show the answer
You are choosing the test set that flatters you, so the number stops meaning anything — Once you choose anything by test score, the test set has quietly joined training.
Classification
Predicting which category something belongs to.