scikit-learn/examples/ensemble/plot_forest_importances.py

56 lines
1.7 KiB
Python
Raw Normal View History

"""
=========================================
Feature importances with forests of trees
=========================================
This examples shows the use of forests of trees to evaluate the importance of
features on an artifical classification task. The red plots are the feature
2012-01-28 21:52:10 +08:00
importances of each individual tree, and the blue plot is the feature
importance of the whole forest.
2012-01-28 21:52:10 +08:00
As expected, the knee in the blue plot suggests that 3 features are
informative, while the remaining are not.
"""
print __doc__
import numpy as np
from sklearn.datasets import make_classification
2011-12-20 02:00:43 +08:00
from sklearn.ensemble import ExtraTreesClassifier
# Build a classification task using 3 informative features
X, y = make_classification(n_samples=1000,
n_features=10,
n_informative=3,
n_redundant=0,
n_repeated=0,
n_classes=2,
random_state=0,
shuffle=False)
# Build a forest and compute the feature importances
2011-12-28 00:54:25 +08:00
forest = ExtraTreesClassifier(n_estimators=250,
compute_importances=True,
random_state=0)
forest.fit(X, y)
2011-12-28 00:54:25 +08:00
importances = forest.feature_importances_
indices = np.argsort(importances)[::-1]
# Print the feature ranking
print "Feature ranking:"
for f in xrange(10):
print "%d. feature %d (%f)" % (f + 1, indices[f], importances[indices[f]])
# Plot the feature importances of the trees and of the forest
import pylab as pl
pl.figure()
pl.title("Feature importances")
for tree in forest.estimators_:
2011-12-28 00:54:25 +08:00
pl.plot(xrange(10), tree.feature_importances_[indices], "r")
pl.plot(xrange(10), importances[indices], "b")
pl.show()