欢迎您访问程序员文章站本站旨在为大家提供分享程序员计算机编程知识!
您现在的位置是: 首页

画出sklearn中的决策树的图

程序员文章站 2024-02-26 15:53:22
...

经常会使用sklearn中的决策树进行机器学习,比如分类,不过很想对其结果进行可视化,话不多说直接上分类树的代码:

import numpy as np
import pandas as pd
from  sklearn.tree import DecisionTreeClassifier
from sklearn.tree import export_graphviz

##准备数据
X=[np.random.rand(5) for i in range(200)]
y=[int(np.random.rand()*5) for i in range(200)]
feature_names=['A','B','C','D','E']
class_names=['qingtong','huangjin','baijin','zuanshi','xingyao']

#训练
tree=DecisionTreeClassifier(max_depth=3,random_state=234)
tree.fit(X,y)
##导出dot文件
export_graphviz(
        tree,
        out_file="C:\\honor_tree.dot",
        feature_names=feature_names,
        class_names=class_names,
        rounded=True,
        filled=True
    )

运行成功后会出现一个dot文件:honor_tree.dot

想要显示出来这个文件里面的图,需要安装graphviz

这个有windows版的,支持很多个系统

我下的是win版的,很小就几兆

安装成功后,就能在命令行下转生成图像了:

C:\>dot -Tpng honor_tree.dot -o honor_tree.png

结果会生成一个名为honor_tree.png的图片,效果如下:

画出sklearn中的决策树的图

 

当然也可以做成回归树

import numpy as np
import pandas as pd
from  sklearn.tree import DecisionTreeRegressor
from sklearn.tree import export_graphviz

##准备数据
X=[np.random.rand(5) for i in range(200)]
y=[np.random.rand()*50 for i in range(200)]
feature_names=['A','B','C','D','E']

#训练
tree=DecisionTreeRegressor(max_depth=3,random_state=234)
tree.fit(X,y)
##导出dot文件
export_graphviz(
        tree,
        out_file="C:\\honor_tree_re.dot",
        feature_names=feature_names,
        rounded=True,
        filled=True
    )

 类似的方式生成图片:

C:\>dot -Tpng honor_tree_re.dot -o honor_tree_re.png

结果如下:

画出sklearn中的决策树的图