Skip to main content
Donate
All projects
Machine learning project

Wine detective with a decision tree

scikit-learn's wine dataset describes 178 wines from three different growers in the same region of Italy, using chemistry measurements like alcohol, colour intensity and flavonoids. A decision tree learns a set of yes/no questions to tell the growers apart, and unlike many models, you can read every question it asks.

Skills you'll practise

  • Decision trees
  • Train/test split
  • Feature importance
  • Overfitting
Open starter in playgroundSee the steps

Steps

  1. Step 1: Load the wines

    Load the dataset as a pandas DataFrame and look at the first few rows and the column names.

    Show a hint

    load_wine(as_frame=True).frame gives you one table with the features and a target column.

  2. Step 2: Split, then train a small tree

    Hold back a quarter of the wines, then fit a DecisionTreeClassifier with max_depth=3 on the rest.

    Show a hint

    A shallow tree is easier to read and less likely to memorise the training data.

  3. Step 3: Test it

    Print the accuracy on the wines the tree hasn't seen.

  4. Step 4: Read the tree's questions

    Print the tree as text with export_text to see each question it asks, in order.

  5. Step 5: Find the most useful clues

    Sort feature_importances_ to see which measurements did most of the work.

    Show a hint

    Put them in a pandas Series with the feature names as the index, then call .sort_values(ascending=False).

Starter code

It already runs. The TODO comments mark where to start.

Open in playground
main.py
import pandas as pd
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text

wine = load_wine(as_frame=True)
df = wine.frame

print(df.head())
print()
print("Growers:", ", ".join(wine.target_names))

# TODO: split into training and test wines, train a
# DecisionTreeClassifier(max_depth=3), and print its accuracy.

Example solution

One way to finish it. Have a go first; yours doesn't need to match.

Reveal the solution
main.py
import pandas as pd
from sklearn.datasets import load_wine
from sklearn.model_selection import train_test_split
from sklearn.tree import DecisionTreeClassifier, export_text

wine = load_wine(as_frame=True)
X = wine.data
y = wine.target

X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, random_state=0, stratify=y
)

tree = DecisionTreeClassifier(max_depth=3, random_state=0)
tree.fit(X_train, y_train)

print(f"Accuracy on unseen wines: {tree.score(X_test, y_test):.0%}")
print()
print("The questions the tree asks:")
print(export_text(tree, feature_names=list(X.columns)))

importance = pd.Series(tree.feature_importances_, index=X.columns)
print("Most useful clues:")
print(importance.sort_values(ascending=False).head(5).round(2))
Open solution in playground

Stretch goal

Train trees with max_depth from 1 to 10 and print both the training and the test accuracy for each. Where does the tree start memorising instead of learning?

Chapters that help