import pydotplus
from IPython.display import Image
from sklearn.tree import export_graphviz
# Graphviz临时环境变量
import os
os.environ['PATH'] += r'.\Graphviz2.38\bin'
def print_graph(dtr, feature_names):
"""绘制决策树"""
graph = export_graphviz(dtr, feature_names=feature_names, class_names={0:"D", 1:"R"},
label="root", proportion=True, impurity=False, out_file=None,
filled=True, rounded=True)
graph = pydotplus.graph_from_dot_data(graph)
return Image(graph.create_png())
# 生成决策树
from sklearn.tree import DecisionTreeClassifier
dtr = DecisionTreeClassifier(max_depth=3, random_state=SEED)
dtr.fit(X, y)
# 绘制
print_graph(dtr, X.columns)
决策树可视化
最新推荐文章于 2024-04-05 22:15:49 发布