网上的一些教程都需要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
