机器学习库的决策树绘制
使用graphviz绘制决策树更好一点,可以避免节点重叠,并且导出的pdf属于非图片型pdf。
绘制决策树大体分为3个步骤:
1.准备数据
2.创建并训练决策树模型
3.设置决策树参数绘制决策树
示例代码如下:
from sklearn import tree
import graphviz#1.准备数据
x=[[2, 0, 0, 2],[2, 1, 1, 0],[2, 1, 1, 2],[2, 1, 1, 1],[2, 0, 1, 0],[2, 0, 1, 1],[2, 0, 0, 1],[0, 1, 1, 0],[0, 1, 1, 1],[0, 1, 0, 0],[0, 1, 0, 1],[0, 0, 1, 2],[0, 0, 1, 1],[1, 1, 0, 0],[1, 1, 0, 2],[1, 0, 1, 0],[1, 0, 1, 2],[1, 0, 1, 1],[1, 0, 0, 0],[1, 0, 0, 1]]
y = [1, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0]#设置类别映射表
class_map = {0: '不满意',1: '满意'
}#2.创建并训练决策树模型
model = tree.DecisionTreeClassifier()
model.fit(x, y)#3.设置决策树参数绘制决策树
dot_data = tree.export_graphviz(model,feature_names=['维护类型','环境温度','环境湿度','通风效能'],class_names = list(map(lambda x:class_map[x],model.classes_)),filled = True,rounded = True,fontname='SimHei',special_characters=True,leaves_parallel=False)
graph = graphviz.Source(dot_data)
graph.render("DecisionTree", format="png")
graph.view()