用Python实现决策树用于离职员工分析详细过程(一)(含首次调试运行)文章中进行了决策树模型分析离职数据的库安装、建模及训练,本节将对训练好的模型开展评估预测。训练好模型后,需要使用测试集对模型进行评估。
1、直接预测是否离职
这里把测试集中的数据导入模型中进行预测,代码如下,其中model就是前面训练好的决策树模型:
y_pred = dt_model.predict(X_test)
print(y_pred)
predict()函数将直接打印预测结果,其中0为预测不离职,1为预测离职。
运行结果,发现如上图所示,在idle中,未完全显示预测结果,那我们如何全部显示完整呢?首先,我们需要知道dt_model.predict(X_test)返回的是什么类型,为此使用如下代码进行判断:
print(type(y_pred))
运行结果如下图所示: