员工离职预测模型 | 更文挑战

751 阅读2分钟

“这是我参与更文挑战的第23天,活动详情查看: 更文挑战

通过已有的员工信息和离职表现来搭建相应的员工离职预测模型,可以预测之后的员工是否会离职。

1.数据读取与预处理

首先读取员工信息和离职表现(即是否离职的记录),代码如下。

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

通过打印df.head()查看表格的前5行,结果如下图所示,"离职"列中的数字1代表离职,0代表未离职。

image.png

该表格共有15000组历史数据,前3571组为离职员工数据,后11429组为非离职员工数据。我们的目的就是根据这些历史数据搭建决策树模型来预测之后员工的离职可能性。

原始数据中的“工资”数据被分为“高”“中”“低”3个等级,而Python数学建模中无法识别这种文本内容,所以“工资”列的内容需要进行数值化处理。这里用pandas库中的replace()函数将文本“高”“中”“低”分别替换为数字2、1、0。

df=df.replace({'工资':{'低'0,'中':1,'高':2}})
df.head()

处理后的表格前5行如下图所示。

image.png

将表格中的“离职”列作为目标变量,剩下的字段作为特征变量,通过一个员工的特征来判断其是否会离职。为便于演示,这里只选取了6个特征变量,在商业实战中选取的特征变量会更多。接下来进行决策树模型的搭建,这是大部分机器学习模型搭建的常规步骤。

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

首先将特征变量和目标变量分别提取出来,代码如下

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

第1行代码用drop()函数删除“离职”列,将剩下的数据作为特征变量赋给变量X;第2行代码用DataFrame提取列的方式提取“离职”列作为目标变量,并赋给变量y。

3.划分训练集和测试集

提取完特征变量和目标变量后,还需要将原来的15000组数据划分为训练集和测试集。顾名思义,训练集用于进行模型训练,而测试集用于检验模型训练的结果。代码如下。

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)

划分后的数据如下图所示。

image.png

image.png

4.模型训练及搭建

划分好训练集和测试集之后,就可以从Scikit-Learn库中引入决策树模型进行模型训练了,代码如下。

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

至此,一个决策树模型便搭建完成了