網上的一些教程都需要graphviz這個模塊,然后dot什么的,走了一大圈彎路也沒有成功.還是官方文檔靠譜,不過官網上的代碼不帶彩色,憑著一頓搗鼓和猜測最終實現(xiàn)了使用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
