分享一个Python可视化工具pybaobabdt,轻松对决策树、随机森林可视化,例如,
图怎么看:每一种颜色代表一个class,link的宽度表示从一个节点流向另一个节点的items数量。
安装
pip install pybaobabdt
pip install pygraphviz
可视化决策树
import pybaobabdt
import pandas as pd
from scipy.io import arff
from sklearn.tree import DecisionTreeClassifier #导入决策树算法
data = arff.loadarff('winequality-red.arff')
df = pd.DataFrame(data[0])
y = list(df['class'])
features = list(df.columns)
features.remove('class')
X = df.loc[:, features]
clf = DecisionTreeClassifier().fit(X, y)
ax = pybaobabdt.drawTree(clf, size=10, dpi=300, features=features) #可视化主函数pybaobabdt.drawTree