基于 Python 的随机森林和 XGBoost 薪资预测模型：从数据爬取到预测优化
引言
在招聘市场中，职位的薪资水平是求职者和企业关注的核心指标之一。然而，由于不同职位的薪资标准受多方面因素影响，例如年限要求、学历要求、公司规模、行业类型等，因此单纯依靠经验进行判断难以精准。为此，我们通过 Python 数据爬取、数据清洗、特征工程、机器学习建模等步骤，构建出薪资预测模型，帮助用户更加直观地了解市场薪酬水平。
系列博客结构
-
Python 爬取 BOSS 直聘职位数据并保存到 Excel： 通过 Python 结合爬虫库，批量获取 BOSS 直聘的职位数据，存储为 Excel 文件格式，为后续分析提供数据支持。
-
Python 清洗与深度解析 BOSS 直聘数据： 对爬取到的数据进行标准化和清洗，包括薪资转换、年限解析、公司规模分析等，以保证模型输入的数据质量。
-
揭秘职场薪资秘密：用 Python 全面解读 BOSS 直聘招聘数据：
基于清洗后的数据展开深入分析，围绕薪资分布、职位要求、行业分布等多个维度，揭示职场薪资秘密、技能需求和地域分布特点
-
Python 实现岗位需求趋势与薪资预测（本文）： 使用随机森林和XGBoost模型对招聘数据进行薪资预测，进行模型性能比较，并通过可视化和解释工具深入分析模型输出结果。
环境配置与库安装
1. 需要安装的 Python 库
本文使用以下 Python 库：
pandas：数据读取与处理。numpy：数值运算。scikit-learn：机器学习模型库。xgboost：提升树模型库。matplotlib和seaborn：可视化工具。shap：解释机器学习模型输出的重要性。joblib：用于保存和加载模型。- nltk：停用词处理
2. 安装指令
建议使用阿里云镜像源加速库安装：
pip install -i https://mirrors.aliyun.com/pypi/simple pandas numpy scikit-learn xgboost matplotlib seaborn shap joblib openpyxl nltk
数据预处理与特征工程
读取 Excel 数据
import pandas as pd
import numpy as np
import re
from sklearn.preprocessing import OneHotEncoder
from sklearn.model_selection import train_test_split
# 读取 Excel 文件
df = pd.read_excel('job_data.xlsx')
# 删除缺失值
df.dropna(subset=['岗位要求', '薪资描述'], inplace=True)
标准化薪资字段
我们将薪资范围（如 8k-15k/月 或 年薪15万·14薪）标准化为中位数值：
def extract_salary(salary_text):
salaries = re.findall(r'\d+', salary_text)
if len(salaries) == 2:
return (int(salaries[0]) + int(salaries[1])) // 2
elif len(salaries) == 1:
return int(salaries[0])
else:
return np.nan
df['平均薪资'] = df['薪资描述'].apply(extract_salary)
df.dropna(subset=['平均薪资'], inplace=True)
转换其他字段
将是否上市转化为二进制值，将公司规模中的人数标准化为数值：
df['是否上市'] = df['是否上市'].apply(lambda x: 1 if '上市' in x else 0)
df['公司规模'] = df['公司规模'].apply(lambda x: int(re.findall(r'\d+', str(x))[0]) if re.findall(r'\d+', str(x)) else 10)
独热编码特征
对城市和行业类型进行独热编码：
encoder = OneHotEncoder(sparse=False, handle_unknown='ignore')
encoded_features = pd.DataFrame(encoder.fit_transform(df[['城市', '行业类型']]), columns=encoder.get_feature_names(['城市', '行业类型']))
df = pd.concat([df, encoded_features], axis=1)
df.drop(columns=['城市', '行业类型'], inplace=True)
数据集划分
features = ['年限要求', '学历要求', '公司规模', '是否上市'] + list(encoded_features.columns)
X = df[features]
y = df['平均薪资']
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
模型构建与调参
1. 随机森林模型
随机森林是一种集成学习方法，通过构建多棵决策树来提高模型预测的稳定性和准确性。其优势在于抗过拟合能力强，适合处理高维数据。
from sklearn.ensemble import RandomForestRegressor
from sklearn.model_selection import GridSearchCV
param_grid_rf = {
'n_estimators': [100, 200],
'max_depth': [10, 20, None],
'min_samples_split': [2, 5]
}
grid_rf = GridSearchCV(RandomForestRegressor(random_state=42), param_grid_rf, cv=3, scoring='neg_mean_absolute_error')
grid_rf.fit(X_train, y_train)
2. XGBoost 模型
XGBoost 是一种基于梯度提升的树模型，具有高效的并行处理能力和强大的非线性拟合能力，适用于大规模数据集。
from xgboost import XGBRegressor
param_grid_xgb = {
'n_estimators': [100, 200],
'learning_rate': [0.01, 0.1],
'max_depth': [3, 6]
}
grid_xgb = GridSearchCV(XGBRegressor(random_state=42), param_grid_xgb, cv=3, scoring='neg_mean_absolute_error')
grid_xgb.fit(X_train, y_train)
模型评估与可视化
1. 模型评估
使用 MAE 和 RMSE 作为评估指标：
from sklearn.metrics import mean_absolute_error, mean_squared_error
def evaluate_model(y_true, y_pred, model_name):
mae = mean_absolute_error(y_true, y_pred)
rmse = np.sqrt(mean_squared_error(y_true, y_pred))
print(f"{model_name} - MAE: {mae:.2f}, RMSE: {rmse:.2f}")
# 评估模型
rf_preds = grid_rf.best_estimator_.predict(X_test)
xgb_preds = grid_xgb.best_estimator_.predict(X_test)
evaluate_model(y_test, rf_preds, '随机森林')
evaluate_model(y_test, xgb_preds, 'XGBoost')
2. 可视化预测结果
import matplotlib.pyplot as plt
plt.figure(figsize=(10, 6))
plt.plot(y_test.values, label='实际薪资', alpha=0.8)
plt.plot(rf_preds, label='随机森林预测薪资', alpha=0.8)
plt.plot(xgb_preds, label='XGBoost预测薪资', alpha=0.8)
plt.legend()
plt.title('实际薪资 vs 预测薪资')
plt.show()
模型解释
使用 SHAP 工具分析模型输出：
import shap
explainer = shap.Explainer(grid_xgb.best_estimator_)
shap_values = explainer(X_test)
shap.summary_plot(shap_values, X_test)
模型保存与加载
import joblib
joblib.dump(grid_xgb.best_estimator_, 'xgb_salary_model.pkl')
loaded_model = joblib.load('xgb_salary_model.pkl')
4. 结果分析与总结
-
模型性能：
- 随机森林和 XGBoost 都表现出较高的薪资预测能力，尤其在参数优化后，误差显著降低。
- 在我们的实验中，XGBoost 的 RMSE 较低，表明其在非线性数据上的拟合能力优于随机森林。
-
特征重要性：
- 工作年限和公司规模是影响薪资的重要特征，而行业类型和是否上市在某些情况下对高薪岗位的影响较大。
-
模型优化方向：
- 可以进一步引入 NLP 模型对岗位描述进行语义分析，提高模型对技能要求的理解。
- 引入时间特征，分析职位薪资随时间的趋势变化。
代码实现
import pandas as pd
import numpy as np
import re
from sklearn.model_selection import train_test_split
from sklearn.metrics import mean_squared_error, mean_absolute_error
from xgboost import XGBRegressor
import optuna
from xgboost.testing.data import joblib
# **1. 数据预处理方法**
def load_and_preprocess_data(file_path):
# 加载 Excel 数据
df = pd.read_excel(file_path)
# 删除缺失值
df.dropna(subset=['岗位要求', '薪资描述'], inplace=True)
# 提取薪资范围的平均值
def extract_salary(salary_text):
salaries = re.findall(r'\d+', str(salary_text))
if len(salaries) == 2:
return (int(salaries[0]) + int(salaries[1])) // 2
elif len(salaries) == 1:
return int(salaries[0])
else:
return np.nan
df['平均薪资'] = df['薪资描述'].apply(extract_salary)
df.dropna(subset=['平均薪资'], inplace=True)
# 将是否上市转换为二进制变量
df['是否上市'] = df['是否上市'].apply(lambda x: 1 if x == '已上市' else 0)
# 将公司规模转换为数值
df['公司规模'] = df['公司规模'].apply(
lambda x: int(re.findall(r'\d+', str(x))[0]) if re.findall(r'\d+', str(x)) else 10
)
# 学历要求数值映射
education_mapping = {'博士': 4, '硕士': 3, '本科': 2, '大专': 1, '不限': 0}
df['学历要求'] = df['学历要求'].map(education_mapping)
# 选择特征和标签
features = ['年限要求', '学历要求', '公司规模', '是否上市']
X = df[features]
y = df['平均薪资']
return X, y
# **2. 使用 Optuna 调参**
def objective(trial, X_train, y_train, X_test, y_test):
# 定义超参数搜索空间
params = {
'n_estimators': trial.suggest_int('n_estimators', 100, 300),
'learning_rate': trial.suggest_float('learning_rate', 0.01, 0.2),
'max_depth': trial.suggest_int('max_depth', 3, 10),
'subsample': trial.suggest_float('subsample', 0.6, 1.0),
'colsample_bytree': trial.suggest_float('colsample_bytree', 0.6, 1.0),
'reg_alpha': trial.suggest_float('reg_alpha', 0.0, 1.0),
'reg_lambda': trial.suggest_float('reg_lambda', 0.0, 1.0),
'random_state': 42
}
# 训练模型
model = XGBRegressor(**params)
model.fit(X_train, y_train)
# 模型预测
y_pred = model.predict(X_test)
# 计算 RMSE 作为目标函数
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
return rmse
def train_model_with_optuna(X, y):
# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
# 创建 Optuna 调参任务
study = optuna.create_study(direction='minimize')
study.optimize(lambda trial: objective(trial, X_train, y_train, X_test, y_test), n_trials=10)
print("最佳参数：", study.best_params)
# 使用最佳参数训练最终模型
best_params = study.best_params
model = XGBRegressor(**best_params, random_state=42)
model.fit(X_train, y_train)
# 模型评估
y_pred = model.predict(X_test)
mae = mean_absolute_error(y_test, y_pred)
rmse = np.sqrt(mean_squared_error(y_test, y_pred))
print(f"最终模型 - MAE: {mae:.2f}, RMSE: {rmse:.2f}")
return model
# **3. 主函数**
def main():
# 文件路径
file_path = 'cleaned_job_data.xlsx'
# 加载数据并预处理
X, y = load_and_preprocess_data(file_path)
# 使用 Optuna 进行调参和训练
print("开始使用 Optuna 调参...")
best_model = train_model_with_optuna(X, y)
# 保存模型
joblib.dump(best_model, 'xgb_salary_model.pkl')
print("模型已保存为 xgb_salary_model.pkl")
# **运行主程序**
if __name__ == "__main__":
main()
代码说明
-
数据预处理：
- 处理缺失值，将薪资描述提取为平均值。
- 将 "是否上市" 转换为二进制值（0 或 1）。
- 将 "公司规模" 转换为数值特征。
- 通过
OneHotEncoder对 "城市" 和 "行业类型" 进行独热编码。
-
模型训练与调参：
- 使用
GridSearchCV对随机森林和 XGBoost 模型进行参数调优。 cv=3表示 3 折交叉验证，选取 MAE（平均绝对误差）作为调参评分标准。
- 使用
-
模型评估：
- 使用 MAE 和 RMSE（均方根误差）评估模型效果。
- 绘制图表对比实际薪资与模型预测薪资。
-
模型解释：
- 使用
SHAP工具展示特征的重要性，分析不同特征对薪资预测的影响。
- 使用
-
模型保存与加载：
- 使用
joblib将训练好的 XGBoost 模型保存到本地文件xgb_salary_model.pkl。 - 通过加载模型进行快速预测验证。
- 使用
总结
-
模型性能表现： 在薪资预测任务中，XGBoost 模型表现优于 随机森林，尤其在处理复杂非线性关系时具有更低的误差，更适用于多维特征和高稀疏性数据。
-
关键特征分析： 数据分析表明，工作年限 和 学历要求 是影响薪资水平的核心因素。此外，是否上市 和 行业类型 等特征在高薪岗位预测中具有显著影响，尤其在技术密集型行业和大规模企业中，薪资水平往往更高。
-
模型优化与扩展方向：
- 文本特征增强：引入 NLP 技术（如 BERT 或 TF-IDF 结合 LDA 主题模型），提取岗位描述中的语义信息，增强模型对技能需求的理解。
- 动态特征引入：可结合时间序列数据，分析不同季节或季度的薪资趋势。
- 多模型融合：探索 LightGBM、CatBoost 等模型，通过加权集成进一步提升预测效果。
希望本文帮助你完整掌握了从数据爬取到薪资预测建模与优化的完整流程，同时为未来的招聘市场薪酬分析提供了更系统的参考方案！
