XGBoost GPU加速实战：优化训练与SHAP值计算的性能考量

from sklearn.datasets import fetch_california_housing
import xgboost as xgb
import time

# 1. 准备数据集
data = fetch_california_housing()
X = data.data
y = data.target

num_round = 1000 # 提升轮数

# 2. CPU多核训练配置
param_cpu = {
    "eta": 0.05,
    "max_depth": 10,
    "tree_method": "hist", # 使用hist方法，可在CPU上高效运行
    "device": "cpu",       # 明确指定使用CPU
    "nthread": 24,         # 根据CPU核心数调整线程数
    "objective": "reg:squarederror",
    "seed": 42
}

# 3. GPU加速训练配置
param_gpu = {
    "eta": 0.05,
    "max_depth": 10,
    "tree_method": "gpu_hist", # 使用gpu_hist方法
    "device": "GPU",           # 明确指定使用GPU
    "objective": "reg:squarederror",
    "seed": 42
}

dtrain = xgb.DMatrix(X, label=y, feature_names=data.feature_names)

print("--- CPU 多核训练开始 ---")
start_time_cpu = time.time()
model_cpu = xgb.train(param_cpu, dtrain, num_round)
end_time_cpu = time.time()
print(f"CPU 训练耗时: {end_time_cpu - start_time_cpu:.2f} 秒")

print("\n--- GPU 加速训练开始 ---")
start_time_gpu = time.time()
model_gpu = xgb.train(param_gpu, dtrain, num_round)
end_time_gpu = time.time()
print(f"GPU 训练耗时: {end_time_gpu - start_time_gpu:.2f} 秒")

import shap

# 确保模型参数已设置为GPU，或者在预测前设置
# model_gpu.set_param({"device": "gpu"}) # 如果模型是在CPU上训练的，需要先切换设备

print("\n--- CPU 计算 SHAP 值开始 ---")
# 默认情况下，predict(pred_contribs=True) 会在CPU上运行，除非模型本身设置为GPU
start_time_shap_cpu = time.time()
# 假设我们用CPU训练的模型来计算SHAP值，或者强制在CPU上计算
shap_values_cpu = model_cpu.predict(dtrain, pred_contribs=True)
end_time_shap_cpu = time.time()
print(f"CPU 计算 SHAP 耗时: {end_time_shap_cpu - start_time_shap_cpu:.2f} 秒")


print("\n--- GPU 加速计算 SHAP 值开始 ---")
# 确保模型已设置为GPU，或者重新加载/设置模型以使用GPU
# 如果model_gpu已经是GPU模型，则无需再次设置
model_gpu.set_param({"device": "GPU"}) # 显式设置，确保使用GPU
start_time_shap_gpu = time.time()
shap_values_gpu = model_gpu.predict(dtrain, pred_contribs=True)
end_time_shap_gpu = time.time()
print(f"GPU 计算 SHAP 耗时: {end_time_shap_gpu - start_time_shap_gpu:.2f} 秒")