如何在SHAP汇总图图例中映射类别ID与类别名称

我在鸢尾花（Iris）数据集上训练了随机森林分类器，代码如下：

from sklearn import datasets
from sklearn.model_selection import train_test_split
from sklearn.ensemble import RandomForestClassifier

iris = datasets.load_iris()
X = iris.data
y = iris.target

# 划分训练集与测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=0)

model = RandomForestClassifier()
model.fit(X_train, y_train)
predictions = model.predict(X_test)

随后通过以下代码绘制SHAP值汇总图：

import shap

explainer = shap.TreeExplainer(model)
shap_values = explainer.shap_values(X_test)

shap.summary_plot(shap_values, X_test)

生成的图中图例显示的是类别ID（Class 0、Class 1、Class 2），需要将其映射为对应的中文类别名称：

• Class 0 → 山鸢尾（对应原名称setosa）
• Class 1 → 变色鸢尾（对应原名称versicolor）
• Class 2 → 弗吉尼亚鸢尾（对应原名称virginica）

数据集自带的类别名称数组为：

iris.target_names
# 输出：array(['setosa', 'versicolor', 'virginica'], dtype='<U10')

实现方法 #

调用shap.summary_plot()时，传入class_names参数指定中文类别名称即可完成替换：

# 定义中文类别名称列表
class_names = ["山鸢尾", "变色鸢尾", "弗吉尼亚鸢尾"]

# 绘制带中文类别的SHAP汇总图
shap.summary_plot(shap_values, X_test, class_names=class_names)

执行上述代码后，汇总图的图例将直接显示中文类别名称，替代原有的Class ID标识。

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