[数据结构与算法]决策树的简单模型

python的基础项目
音乐推送的决策树算法

数据集

一、建立决策树模型

import pandas as pd 
from sklearn.tree import DecisionTreeClassifier   # 决策树分类器
from sklearn.model_selection import train_test_split  #划分数据集函数
from sklearn.metrics import accuracy_score
#from sklearn.externals import joblib
import joblib

music_data = pd.read_csv('music_data.csv')
X = music_data.drop(columns=['genre'])
y = music_data['genre']
X_train,X_test,y_train,y_test = train_test_split(X ,y , test_size = 0.2)
model = DecisionTreeClassifier()
model.fit(X_train,y_train)

predictions = model.predict(X_test)
predictions
score = accuracy_score(y_test, predictions)
score

二、保存模型&调用模型

建模后保存模型

import pandas as pd 
from sklearn.tree import DecisionTreeClassifier   # 决策树分类器
from sklearn.model_selection import train_test_split  #划分数据集函数
from sklearn.metrics import accuracy_score
#from sklearn.externals import joblib
import joblib

music_data = pd.read_csv('music_data.csv')
X = music_data.drop(columns=['genre'])
y = music_data['genre']
X_train,X_test,y_train,y_test = train_test_split(X ,y , test_size = 0.2)
model = DecisionTreeClassifier()
model.fit(X_train,y_train)
#保存模型
joblib.dump(model, "music_recommender.joblib")

加载模型进行预测

#加载模型
model = joblib.load("music_recommender.joblib")
predictions = model.predict(X_test)
predictions
score = accuracy_score(y_test, predictions)
score

三、可视化决策树

import pandas as pd 
from sklearn.tree import DecisionTreeClassifier   # 决策树分类器
from sklearn.model_selection import train_test_split  #划分数据集函数
from sklearn import tree

music_data = pd.read_csv('music_data.csv')
X = music_data.drop(columns=['genre'])
y = music_data['genre']
X_train,X_test,y_train,y_test = train_test_split(X ,y , test_size = 0.2)
model = DecisionTreeClassifier()
model.fit(X_train,y_train)

#输出dot文件
tree.export_graphviz(model, out_file="music_recommender.dot",
                    feature_names = ['age','gender'],
                    class_names = sorted(y.unique()),
                     label = 'all',
                     rounded= True,
                     filled= True)
import graphviz

with open("music_recommender.dot") as f:
    dot_graph = f.read()
dot=graphviz.Source(dot_graph)
dot.view()

注意：anaconda需要安装 python-graphviz和 graphviz两个包