Splitting the data in Python

Never score a model on rows it trained on: split first, always.

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.25, random_state=0, stratify=y)

print(X_train.shape, X_test.shape)
print(y_train[:5], y_test[:5])

How it works

1. test_size is the share held back for the final score.
2. random_state makes the split reproducible.
3. stratify=y keeps the class balance in both halves.

Keywords and builtins used here

• print

The run, in numbers

Lines

9

Characters to type

296

Tokens

72

Three-star pace

100 tpm

At the three-star pace of 100 tokens a minute, this run takes about 43 seconds.

Type this snippet

Step 2 of 3 in The estimator API, step 2 of 25 in Machine learning with scikit-learn.

← Previous Next →