分享一个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

可视化决策树-设置树深度

通过maxdepth限制需要渲染的树深度，对large trees特别友好，

ax = pybaobabdt.drawTree(clf,size=10,dpi=100,maxdepth=6,  #设置渲染的树的最大深度features=features)

可视化决策树-突出class 

通过colormap凸显感兴趣的class，

from matplotlib.colors import ListedColormapax = pybaobabdt.drawTree(clf,size=10,dpi=600,colormap=ListedColormap(["#01a2d9", "gray", "#d5695d","gray"]),  #通过颜色突出展示class 5和6 变量features=features)

可视化随机森林

import pybaobabdt
import pandas as pd
from scipy.io import arff
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier  #导入随机森林算法data = arff.loadarff('vehicle.arff')df = pd.DataFrame(data[0])
y = list(df['class'])
features = list(df.columns)
features.remove('class')
X = df.loc[:, features]clf = RandomForestClassifier(n_estimators=20, n_jobs=-1, random_state=0)
clf.fit(X, y)size = (15, 15)
plt.rcParams['figure.figsize'] = size
fig = plt.figure(figsize=size, dpi=300)for idx, tree in enumerate(clf.estimators_):ax1 = fig.add_subplot(5, 4, idx + 1)pybaobabdt.drawTree(tree,model=clf,size=15,dpi=300,maxdepth=8,features=features,ax=ax1)

参考：https://pypi.org/project/pybaobabdt/