经常会使用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
)
类似的方式生成图片:
C:\>dot -Tpng honor_tree_re.dot -o honor_tree_re.png
结果如下: