scikit-learn
decision tree
machine learning
data science
Python

How do I find which attributes my tree splits on, when using scikit-learn?

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

When using decision trees for machine learning in scikit-learn, understanding the attributes or features on which the tree splits can provide valuable insights into the model's decision-making process. Determining these split features can help you interpret your model, diagnose overfitting, and improve feature engineering.

Decision Trees in Scikit-Learn

A decision tree is a popular supervised machine learning algorithm used for classification and regression tasks. It works by recursively splitting the dataset into subsets based on feature values, resulting in a tree-like model of decisions. In scikit-learn, decision trees for classification tasks are implemented using DecisionTreeClassifier, and for regression tasks, DecisionTreeRegressor is used.

How Decision Trees Split

In decision trees, the dataset is split through a series of questions or decisions. Each internal node represents a feature-based condition, and an edge represents the result of splitting. The leaves represent class labels or target value predictions. The feature selected for splitting at each node is based on a metric like Gini impurity (for classification) or variance reduction (for regression), which aims to maximize the purity or homogeneity of the resultant child nodes.

Finding Split Attributes

The following steps guide you through finding split attributes in a decision tree built using scikit-learn:

  1. Train the Model: First, train your decision tree model using your dataset.
python
1    from sklearn.tree import DecisionTreeClassifier
2    from sklearn.datasets import load_iris
3
4    # Load dataset
5    iris = load_iris()
6    X, y = iris.data, iris.target
7
8    # Train decision tree classifier
9    clf = DecisionTreeClassifier()
10    clf.fit(X, y)
  1. Extract Feature Importances: Scikit-learn provides an attribute feature_importances_, which is a normalized version of the total reduction of the criterion brought by that feature.
python
    feature_importances = clf.feature_importances_
    print("Feature importances:\n", feature_importances)
  1. Use Tree’s Attributes: Access the tree's attributes directly to understand exactly the features involved at each split:
    • tree_.feature: An array where each element is the index of the feature used for the split at that node, or -2 if the node is a leaf.
    • tree_.threshold: An array of thresholds for each split.
python
    tree = clf.tree_
    print("Split features:\n", tree.feature)
    print("Thresholds:\n", tree.threshold)
  1. Mapping to Feature Names: Optionally, map feature indices to names for better interpretability:
python
    feature_names = iris.feature_names
    split_features_names = [feature_names[i] for i in tree.feature if i != -2]
    print("Features used for splitting:\n", split_features_names)

Example Output

After acquiring the feature indices and thresholds, interpreting this information helps to reconstruct how the tree arrived at its decisions, thus providing transparency in the modeling process. Here’s how the output might look:

python
1# Example Outputs
2Feature importances:
3 [0. , 0.0125026, 0.53835801, 0.44913939]
4Split features:
5 [2 3 3 -2 -2 2 -2 -2 2 3 -2 -2]
6Thresholds:
7 [2.45 1.75 1.55 -2.   -2.   4.95 -2.   -2.   4.85 1.65 -2.   -2.  ]
8Features used for splitting:
9 ['petal length (cm)', 'petal width (cm)', 'petal width (cm)', 
10  'petal length (cm)', 'petal length (cm)', 'petal width (cm)']

Key Points Summary

ElementPurpose/Description
feature_importances_Measure of feature's contribution to the model - Reports overall importance, not specific splits
tree_.featureIndex of features used for specific splits - -2 indicates the node is a leaf
tree_.thresholdValues at which to split features at each node
Feature MappingConverts indices to human-readable feature names to understand splits better

Additional Details

Handling Trees for Regression

When dealing with regression trees using DecisionTreeRegressor, the process of identifying split attributes remains largely the same. The primary difference lies in the criterion used, typically the mean squared error, to identify optimal splits.

Visualizing the Tree

Another method to see which attributes are being split is to visualize the tree itself. Scikit-learn provides plot_tree and export_text functions to achieve this.

python
1from sklearn.tree import plot_tree
2import matplotlib.pyplot as plt
3
4plot_tree(clf, feature_names=iris.feature_names, class_names=iris.target_names, filled=True)
5plt.show()

Conclusion

Understanding which features are used in decision tree splits with scikit-learn can greatly inform model interpretation as well as feature engineering practices. By applying the methods described, you can transparently dissect model behavior and drive improvements to model performance and generalization.


Course illustration
Course illustration

All Rights Reserved.