[PYTHON] The tree.plot_tree of scikit-learn was very easy and convenient, so I tried to summarize how to use it easily.
The machine learning library scikit-learn implemented in Python is often used because it makes it easy to experiment with various algorithms. .. Speaking of flower shapes, TensorFlow and PyTorch are hard to use in a rigid field. .. .. With such scikit-learn, a function convenient for drawing after learning "decision tree", which is a typical method of supervised learning, has been implemented from Version 0.21.x, so I tried it while comparing it with the conventional method using GraphViz. It was.
Traditional visualization method: Using GraphViz
Previously, I installed and used another library called GraphViz. It takes a lot of time and effort. .. ..
Install GraphViz@Mac
brew install graphviz
pip install graphviz
Install GraphViz@Ubuntu
sudo apt install -y graphviz
pip install graphviz
Method using GraphViz
import graphviz
from sklearn import tree
iris = load_iris()
clf = DecisionTreeClassifier()
clf = clf.fit(iris.data, iris.target)
graph = graphviz.Source(tree.export_graphviz(clf, class_names=iris.feature_names, filled=True))
graph
Execution result
The execution result can be saved as PDF by executing graph.render ('decision_tree').
Use tree.plot_tree
Let's draw a figure similar to the one drawn using GraphViz using tree.plot_tree. Since it is stored in the tree module of scikit-learn, no additional installation is required.
tree.plot_Method using tree
from sklearn.datasets import load_iris
from sklearn.tree import DecisionTreeClassifier, plot_tree
import matplotlib.pyplot as plt
iris = load_iris()
clf = DecisionTreeClassifier()
clf = clf.fit(iris.data, iris.target)
iris = load_iris()
plt.figure(figsize=(15, 10))
plot_tree(clf, feature_names=iris.feature_names, filled=True)
plt.show()
Execution result
I was able to output the same figure as the method using GraphViz. If you execute it on Jupyter Notebook, you can right-click the drawing result as it is and save it as an image.
Summary
Using scikit-learn's tree.plot_tree and traditional GraphViz method for visualization of decision trees, I found tree.plot_tree easier and more convenient (than traditional methods). I would like to actively utilize it in the future.
Reference
Recommended Posts