决策树绘制决策图 (极简教程)

网上的一些教程都需要graphviz这个模块,然后dot什么的,走了一大圈弯路也没有成功.还是官方文档靠谱,不过官网上的代码不带彩色,凭着一顿捣鼓和猜测最终实现了使用graphviz和dot一样带彩色的决策图. 代码很简单,也不需要安装额外的其他库,只需要matplotlib和sklearn这两个常用的库就可以了.代码如下:

%matplotlib inline
import matplotlib.pyplot as plt
from sklearn.datasets import load_iris
from sklearn import tree

iris = load_iris()
clf = tree.DecisionTreeClassifier()
clf.fit(iris.data, iris.target)

fig = plt.figure(figsize=(12, 12),dpi=100)
tree.plot_tree(clf,  
               feature_names=iris.feature_names,
               class_names=iris.target_names,  
               filled=True, 
               rounded=True,  
              fontsize=None) 
plt.savefig('DTC.png')

效果如下:

决策图.png
官方文档:https://scikit-learn.org/stable/modules/tree.html#tree