二分类任务中生成SHAP汇总图时数组形状错误问题

我为随机森林分类模型生成SHAP汇总图时遇到问题：构建二分类示例数据集并训练模型后，用SHAP TreeExplainer计算测试集的SHAP值形状为(125, 20, 2)，符合多分类输出形式，而非预期的二分类单输出。用该SHAP值生成的bar类型汇总图呈现异常的2×2样式，推测和模型输出类别概率有关，寻求原因及解决方法。

环境信息 #

shap version: 0.45.0
Python version: 3.10.12

代码示例 #

数据生成与模型训练 #

import shap
import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier
from sklearn.datasets import make_classification  # 补充缺失的导入

# 生成合成数据
X, y = make_classification(n_samples=500, n_features=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.25, random_state=42)

# 训练随机森林模型
model = RandomForestClassifier(random_state=42)
model.fit(X_train, y_train)

# 创建SHAP TreeExplainer
explainer = shap.TreeExplainer(model)

# 计算测试集SHAP值
shap_values = explainer.shap_values(X_test)

形状输出 #

X_test shape: (125,20)
shap_values shape: (125, 20, 2)

绘图代码 #

shap.summary_plot(shap_values, X_test, plot_type="bar", max_display=None)

生成的异常图表为2×2布局的bar图，每个子图对应一个类别的特征重要性。

原因分析 #

SHAP的TreeExplainer针对分类模型，默认会返回每个类别对应的SHAP值。对于二分类任务，模型有两个类别（0和1），因此SHAP值形状为(样本数, 特征数, 类别数)即(125,20,2)。当直接将完整的SHAP值传入summary_plot时，函数会为每个类别单独绘制特征重要性bar图，最终呈现2×2的异常布局。

解决方法 #

针对二分类任务，我们通常关注**正类（类别1）**的SHAP值，只需提取对应类别的SHAP值传入绘图函数即可：

方法1：提取正类SHAP值绘图 #

如果shap_values是三维数组，提取类别1的SHAP值：

# 提取类别1的SHAP值，形状为(125,20)
shap_values_positive = shap_values[:, :, 1]
# 生成正常的bar汇总图
shap.summary_plot(shap_values_positive, X_test, plot_type="bar", max_display=None)

如果shap_values是列表（部分SHAP版本中，sklearn分类模型返回[类别0的SHAP值, 类别1的SHAP值]），则取列表第二个元素：

shap_values_positive = shap_values[1]
shap.summary_plot(shap_values_positive, X_test, plot_type="bar", max_display=None)

方法2：使用explainer直接获取预测类别的SHAP值 #

调用explainer(X_test)会返回模型预测类别的SHAP值（形状为(125,20)），直接用于绘图：

shap_values_pred = explainer(X_test)
shap.summary_plot(shap_values_pred.values, X_test, plot_type="bar", max_display=None)

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