小白的机器学习记录——决策树

站在巨人的肩膀上写代码
数据集
链接: https://pan.baidu.com/s/14IsCjgM8Mi-KYrijUaDNsQ 提取码: 2ajn

# -*-coding = utf-8-*-
# @Time:2021/10/1 9:35
# @Author TG
# @File :分类西瓜数据集.py
# @Software: PyCharm
from sklearn.preprocessing import LabelEncoder
from sklearn import tree
import pandas as pd
import joblib
import graphviz

df=pd.read_excel('西瓜数据集.xlsx',engine='openpyxl')

df1=df.iloc[[i for i in range(17)],[1,2,3,4,5,6]]
melon_label=list(df.好瓜)
le = LabelEncoder()														#创建LabelEncoder()对象，用于序列化

for col in df1.columns:											#序列化
    df1[col] = le.fit_transform(df1[col])

clf = tree.DecisionTreeClassifier()						#创建DecisionTreeClassifier()类
clf = clf.fit(df1.values.tolist(), melon_label)     #list型的数字组，label

joblib.dump(clf, "西瓜分类.m")#存储模型，需要使用时直接调用




dot_data = tree.export_graphviz(
    clf,
    feature_names=df1.columns,#特征名称
    class_names=['是','否'],#分类
    filled=True,  # 是否填充颜色，true 填充颜色
    rounded=True,  # 框的形状
)
graph = graphviz.Source(dot_data.replace('helvetica', '"Microsoft YaHei"'), encoding='utf-8')#解决中文不显示
graph.render('西瓜决策树')#生成一个file和一个pdf


print(clf.predict(df1.values.tolist()))         #检验模型正确性

关键函数
joblib：保存模型
tree：生成决策树
graphviz：决策树可视化