Skip to content

scikit-learn train_test_split and your first model

October 29, 2026 · 2 min read

Call train_test_split to hold back part of your labeled data, fit a model on the rest, then score it on the held-back part. Compare that score with a dummy baseline so you know whether the model learned anything.

Split, fit, score

from sklearn.datasets import load_breast_cancer
from sklearn.dummy import DummyClassifier
from sklearn.model_selection import train_test_split
from sklearn.neighbors import KNeighborsClassifier

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.25, stratify=y, random_state=42)
print(len(X_train), len(X_test))

knn = KNeighborsClassifier(n_neighbors=5).fit(X_train, y_train)
dummy = DummyClassifier(strategy="most_frequent").fit(X_train, y_train)
print("knn:  ", round(knn.score(X_test, y_test), 3))
print("dummy:", round(dummy.score(X_test, y_test), 3))

Output:

426 143
knn:   0.93
dummy: 0.629

What each part does

  • X holds the features (a table) and y the target. train_test_split returns four pieces, always in the order X_train, X_test, y_train, y_test.
  • test_size=0.25 keeps a quarter of the rows hidden. The model never sees them while training.
  • stratify=y keeps the class mix the same in both parts, which matters when one class is rare.
  • random_state=42 fixes the shuffle, so you get the same split every run.
  • KNeighborsClassifier predicts by letting the 5 most similar training rows vote. fit stores the training data, and score returns accuracy on the rows you give it.
  • DummyClassifier(strategy="most_frequent") always predicts the most common class. That is your baseline: a model that learns nothing.

The kNN model gets 93% of the test rows right, against 62.9% for the dummy. That gap tells you the model found a real pattern.

Common mistake

Scoring on the training data. knn.score(X_train, y_train) looks great because the model has already seen those rows. Always score on X_test, y_test.

Skipping the baseline. If 90% of your rows belong to one class, a model that always guesses that class scores 90%. Check the dummy first, and only be impressed by a model that clearly beats it.

Takeaway

Split first, fit on the training part, score on the test part, and always compare with a DummyClassifier. A score only means something next to a baseline.

More on Machine Learning