决策树可视化

韭浪

于 2020-07-12 23:42:45 发布

阅读量213

点赞数

CC 4.0 BY-SA版权

文章标签：可视化决策树

原文链接：https://scikit-learn.org

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)