调用scikitplot的plot_roc_curve出现IndexError数组维度错误求解答

报错根因 #

scikitplot的plot_roc_curve接口要求传入的预测概率参数y_probas必须为二维数组，维度格式为(样本数, 类别数)，每一列对应对应类别的预测概率。你代码存在两处问题触发报错：

• 错误调用model.predict()获取结果，该方法输出的是预测类别标签，不是ROC曲线需要的类别概率
• 使用np.squeeze将原本的二维概率数组压缩为(982,)的一维数组，库源码尝试用[:, i]索引第二维时触发维度不匹配错误

解决方法 #

方案1：继续使用scikitplot接口 #

删掉np.squeeze调用，替换predict为predict_proba获取原始二维概率数组：

import scikitplot as skplt
# 替换predict为predict_proba，去掉np.squeeze，保留(982, 2)的二维输出
Y_test_probs = model_cEXT.predict_proba(X_test)
skplt.metrics.plot_roc_curve(y_test['cEXT'], Y_test_probs,
                       title="Digits ROC Curve", figsize=(12,6))

方案2：仅绘制二分类正类ROC（用sklearn原生接口） #

如果不需要同时展示所有类别的ROC曲线，也可以直接调用sklearn原生接口实现：

from sklearn.metrics import roc_curve, auc
import matplotlib.pyplot as plt

# 取正类对应的预测概率列
Y_test_pos_proba = model_cEXT.predict_proba(X_test)[:, 1]
fpr, tpr, _ = roc_curve(y_test['cEXT'], Y_test_pos_proba)
roc_auc = auc(fpr, tpr)

plt.figure(figsize=(12,6))
plt.plot(fpr, tpr, label=f'ROC曲线 (AUC值 = {roc_auc:.2f})')
plt.plot([0, 1], [0, 1], 'r--', label='随机猜测基准线')
plt.xlim([0.0, 1.0])
plt.ylim([0.0, 1.05])
plt.xlabel('假阳性率')
plt.ylabel('真阳性率')
plt.title('Digits ROC Curve')
plt.legend(loc="lower right")
plt.show()

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