决策树——员工离职预测模型搭建

1.数据读取与预处理

import pandas as pd
df = pd.read_excel('员工离职预测模型.xlsx')
df.head()

#处理文本内容，代码如下：
df = df.replace({'工资': {'低': 0, '中': 1, '高': 2}})
df.head()

 

 2.提取特征变量和目标变量

X = df.drop(columns='离职') 
y = df['离职']

3.划分训练集和测试集

from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=123)

4.模型训练及搭建

from sklearn.tree import DecisionTreeClassifier
model = DecisionTreeClassifier(max_depth=3, random_state=123) 
model.fit(X_train, y_train)

上述模型搭建代码汇总

# 1.读取数据与简单预处理
import pandas as pd
df = pd.read_excel('员工离职预测模型.xlsx')
df = df.replace({'工资': {'低': 0, '中': 1, '高': 2}})

# 2.提取特征变量和目标变量
X = df.drop(columns='离职') 
y = df['离职']   

# 3.划分训练集和测试集
from sklearn.model_selection import train_test_split
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=123)

# 4.模型训练及搭建
from sklearn.tree import DecisionTreeClassifier
model = DecisionTreeClassifier(max_depth=3, random_state=123)
model.fit(X_train, y_train)    

模型预测及评估

1.直接预测是否离职

y_pred = model.predict(X_test)
print(y_pred[0:100])

# 通过构造DataFrame进行对比
a = pd.DataFrame()  # 创建一个空DataFrame 
a['预测值'] = list(y_pred)
a['实际值'] = list(y_test)
a.head()

 

# 如果要查看整体的预测准确度，可以采用如下代码：
from sklearn.metrics import accuracy_score
score = accuracy_score(y_pred, y_test)
print(score)

 

# 或者用模型自带的score函数查看预测准确度
model.score(X_test, y_test)

 

2.预测不离职&离职概率 

#通过如下代码查看预测属于各个分类的概率
y_pred_proba = model.predict_proba(X_test)
print(y_pred_proba[0:5])

b = pd.DataFrame(y_pred_proba, columns=['不离职概率', '离职概率']) 
b.head()

 

3.模型预测效果评估 

#求出在不同阈值下的命中率（TPR）以及假警报率（FPR）的值，从而可以绘制ROC曲线
from sklearn.metrics import roc_curve
fpr, tpr, thres = roc_curve(y_test, y_pred_proba[:,1])

#通过相关代码可以查看不同阈值下的假警报率和命中率，代码如下：
a = pd.DataFrame()  # 创建一个空DataFrame 
a['阈值'] = list(thres)
a['假警报率'] = list(fpr)
a['命中率'] = list(tpr)
a

已知了不同阈值下的假警报率和命中率，可绘制ROC曲线，代码如下：
import matplotlib.pyplot as plt
plt.plot(fpr, tpr)
plt.show()

 

#通过如下代码则可以快速求出模型的AUC值：
from sklearn.metrics import roc_auc_score
score = roc_auc_score(y_test, y_pred_proba[:,1])
print(score)

 

4.特征重要性评估

# 通过DataFrame进行展示，并根据重要性进行倒序排列
features = X.columns  # 获取特征名称
importances = model.feature_importances_  # 获取特征重要性
# 通过二维表格形式显示
importances_df = pd.DataFrame()
importances_df['特征名称'] = features
importances_df['特征重要性'] = importances
importances_df.sort_values('特征重要性', ascending=False)