Chapters
Python114 chapters

Machine learningChapter 102 of 114

Your First Model

Load data, train, predict and score, in fifteen lines.

The whole thing, once

Every scikit-learn model works the same way: build it, fit it on training data, then predict. Here is a complete program that learns to tell three species of iris apart from four measurements.

Python needs scikit-learn, downloaded on first run
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

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
)

model = DecisionTreeClassifier(random_state=0)
model.fit(X_train, y_train)

print("trained on", X_train.shape[0], "examples")
print("accuracy on unseen data:", round(model.score(X_test, y_test), 2))

Output

trained on 120 examples
accuracy on unseen data: 1.0

That is a working classifier, and it got all thirty unseen flowers right.

Be suspicious of that rather than pleased. Iris is a small, clean, famously separable dataset, and thirty test flowers is not many. Real data does not behave like this, and a perfect score is more often a sign that something is wrong than that something is right. The rest of this section is about understanding each of those lines and knowing when to distrust that number.

What the data looks like

Python needs scikit-learn, downloaded on first run
from sklearn.datasets import load_iris

data = load_iris()

print(data.data.shape)
print(data.feature_names)
print([str(name) for name in data.target_names])
print(data.data[0])
print(data.target[:5])

Output

(150, 4)
['sepal length (cm)', 'sepal width (cm)', 'petal length (cm)', 'petal width (cm)']
['setosa', 'versicolor', 'virginica']
[5.1 3.5 1.4 0.2]
[0 0 0 0 0]

150 flowers, 4 measurements each. The target is 0, 1 or 2 — models work in numbers, and target_names maps them back to something readable.

Predicting

predict takes the same shape of input and returns labels:

Python needs scikit-learn, downloaded on first run
from sklearn.datasets import load_iris
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier

data = load_iris()
X_train, X_test, y_train, y_test = train_test_split(
    data.data, data.target, test_size=0.2, random_state=0
)

model = DecisionTreeClassifier(random_state=0).fit(X_train, y_train)

predictions = model.predict(X_test[:5])
print(predictions)
print([str(data.target_names[p]) for p in predictions])
print("actual:", [str(data.target_names[a]) for a in y_test[:5]])

Output

[2 1 0 2 0]
['virginica', 'versicolor', 'setosa', 'virginica', 'setosa']
actual: ['virginica', 'versicolor', 'setosa', 'virginica', 'setosa']

fit returns the model, which is why the chained .fit(...) on one line works.

One new flower

Predicting on a single example needs a 2D input — a table with one row, not a bare list:

Python needs scikit-learn, downloaded on first run
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier

data = load_iris()
model = DecisionTreeClassifier(random_state=0).fit(data.data, data.target)

flower = [[5.1, 3.5, 1.4, 0.2]]
print(data.target_names[model.predict(flower)[0]])

Output

setosa

How confident is it

Most classifiers can give probabilities as well as an answer:

Python needs scikit-learn, downloaded on first run
from sklearn.datasets import load_iris
from sklearn.linear_model import LogisticRegression

data = load_iris()
model = LogisticRegression(max_iter=1000).fit(data.data, data.target)

chances = model.predict_proba([[6.0, 3.0, 4.8, 1.8]])[0]
for name, chance in zip(data.target_names, chances):
    print(f"{str(name):12} {chance:.2f}")

Output

setosa       0.00
versicolor   0.44
virginica    0.56

A prediction with a probability is far more useful than one without. It lets you say "not sure, ask a person" instead of guessing. This flower is a coin toss: the model answers virginica, and it barely believes it.

The pattern to remember

text
X, y = the data
X_train, X_test, y_train, y_test = train_test_split(X, y)
model = SomeModel()
model.fit(X_train, y_train)
model.score(X_test, y_test)
model.predict(new_data)

Swapping DecisionTreeClassifier for almost any other scikit-learn model changes nothing else. That consistency is the library's best feature.

Test yourself

2 questions

What are the three calls every scikit-learn model shares?

Show the answer

fit, predict and score — Swapping one model for another changes nothing else, which is the library's best feature.

Why does model.predict([5.1, 3.5, 1.4, 0.2]) fail?

Show the answer

predict wants rows, so a single example still needs outer brackets — The error says 'Expected 2D array, got 1D array instead' and the fix is another pair of brackets.

Next chapter

Features and Labels

Shaping your data into the grid of numbers a model expects.