画出sklearn中的决策树的图

经常会使用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的图片，效果如下：

当然也可以做成回归树：

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
    )

类似的方式生成图片：

扫描二维码关注公众号，回复： 12062238 查看本文章

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

结果如下：

画出sklearn中的决策树的图

猜你喜欢