如何绘制决策树GridSearchCV中AUC分数的参数热力图

核心思路 #

GridSearchCV的cv_results_字典包含了所有参数组合及对应分数，只需将这些数据转换为透视表格式（行对应max_depth，列对应min_samples_split，单元格值为AUC分数），即可直接用于热力图绘制。

完整实现代码 #

# 导入所需库
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt

# 从cv_results_中提取参数与AUC分数，转为DataFrame
results_df = pd.DataFrame(clf.cv_results_)
# 筛选需要的字段：max_depth、min_samples_split、测试集AUC（若用训练集则替换为mean_train_score）
plot_data = results_df[['param_max_depth', 'param_min_samples_split', 'mean_test_score']]

# 转换为透视表，适配热力图格式
heatmap_data = plot_data.pivot(
    index='param_max_depth',
    columns='param_min_samples_split',
    values='mean_test_score'
)

# 绘制热力图
sns.set_style("whitegrid")
plt.figure(figsize=(8, 6))
ax = sns.heatmap(heatmap_data, annot=True, fmt=".3f", cmap="YlGnBu", cbar=True)

# 设置轴标签与标题
ax.set_xlabel('min_samples_split')
ax.set_ylabel('max_depth')
ax.set_title('GridSearchCV Test AUC Heatmap')

plt.show()

关键说明 #

1. 数据转换：pivot方法将长格式的参数-分数数据转为宽格式，每个单元格对应一组参数组合的AUC结果；
2. 热力图配置：annot=True显示单元格内的AUC数值，fmt=".3f"控制保留三位小数，cmap可根据需求更换颜色主题（如coolwarm）；
3. 切换数据集：若需绘制训练集AUC热力图，只需将mean_test_score替换为mean_train_score即可。

内容的提问来源于stack exchange，提问作者immrm