中文员工离职随机森林可视化代码
时间: 2023-09-23 14:11:32 浏览: 410
以下是一个简单的中文员工离职预测的随机森林可视化代码示例:
```python
# 导入必要的库
import pandas as pd
from sklearn.ensemble import RandomForestClassifier
from sklearn.tree import export_graphviz
import pydot
# 加载数据
data = pd.read_csv('employee_turnover.csv', header=0)
# 将分类变量转换为数值变量
data['salary'] = data['salary'].map({'low': 0, 'medium': 1, 'high': 2})
data['dept'] = pd.factorize(data['dept'])[0]
# 分割数据集为训练集和测试集
train = data.sample(frac=0.8, random_state=1)
test = data.drop(train.index)
# 定义自变量和因变量
features = ['satisfaction_level', 'last_evaluation', 'number_project', 'average_montly_hours', 'time_spend_company', 'Work_accident', 'promotion_last_5years', 'salary', 'dept']
target = 'left'
# 训练随机森林模型
rf = RandomForestClassifier(n_estimators=100, max_features=3, oob_score=True)
rf.fit(train[features], train[target])
# 输出模型精度
print('Out-of-bag score estimate:', rf.oob_score_)
print('Test data accuracy:', rf.score(test[features], test[target]))
# 可视化决策树
tree = rf.estimators_[5]
export_graphviz(tree, out_file='tree.dot', feature_names=features, class_names=['0', '1'], rounded=True, proportion=False, precision=2, filled=True)
(graph, ) = pydot.graph_from_dot_file('tree.dot')
graph.write_png('tree.png')
```
这个代码将员工满意度、最近一次评估、参与的项目数量、平均每月工作小时数、在公司工作时长、是否发生工伤、过去5年是否晋升、薪资和部门等变量用于预测员工是否会离职。随机森林模型使用100棵树,每棵树最多使用3个特征进行拟合。最后,它将第6棵树可视化为图像(tree.png)。
阅读全文