如何在scikit-learn中生成无额外列的4×2混淆矩阵？

我有一个二分类器，训练用的是可信数据，测试时同时用了可信数据和真实场景下的非优质数据。现在生成混淆矩阵时出现了不必要的列，我期望得到一个4×2的矩阵，请问该怎么实现？以下是我的测试代码：

import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

y_test = [0, 1, 1, 1, 2, 2, 3]
predictions = [0, 1, 1, 1, 0, 1, 0]

cm = confusion_matrix(y_test, predictions)

l = ["M", "F", "M?", "F?"]
disp = ConfusionMatrixDisplay(confusion_matrix=cm, display_labels=l)
disp.plot()
plt.show()
# I expect a 4x2 matrix here. How can I do it?

问题根源 #

你的y_test包含4个类别（0、1、2、3），但predictions只有2个类别（0、1），sklearn默认的confusion_matrix会把所有出现过的类别都纳入矩阵维度，所以生成了4×4的矩阵，这就是你看到额外列的原因。

解决方法 #

要生成4×2的混淆矩阵，需要明确指定行对应真实标签的4个类别，列对应预测的2个类别，可以通过以下两种方式实现：

方法1：映射真实标签到二分类逻辑后生成矩阵

既然是二分类器，M?和F?本质上对应二分类的两个类别（比如M?对应0，F?对应1），先把y_test里的2、3映射到0、1，再计算混淆矩阵：

import matplotlib.pyplot as plt
from sklearn.metrics import confusion_matrix, ConfusionMatrixDisplay

y_test = [0, 1, 1, 1, 2, 2, 3]
predictions = [0, 1, 1, 1, 0, 1, 0]

# 把真实标签映射为二分类：2→0，3→1，0和1保持不变
y_test_binary = [0 if x in [0,2] else 1 for x in y_test]

cm = confusion_matrix(y_test_binary, predictions)
# 行显示原4类标签，列显示二分类标签
disp = ConfusionMatrixDisplay(
    confusion_matrix=cm,
    display_labels=[["M", "F", "M?", "F?"], ["M", "F"]]
)
disp.plot()
plt.show()

方法2：手动构造4×2的混淆矩阵

如果需要保留真实标签的4个类别作为行维度，直接统计每个真实类别被预测为0和1的样本数：

import matplotlib.pyplot as plt
import numpy as np
from sklearn.metrics import ConfusionMatrixDisplay

y_test = [0, 1, 1, 1, 2, 2, 3]
predictions = [0, 1, 1, 1, 0, 1, 0]

# 初始化4行2列的空矩阵
cm = np.zeros((4, 2), dtype=int)
# 遍历样本统计对应类别的数量
for true_label, pred_label in zip(y_test, predictions):
    cm[true_label][pred_label] += 1

# 指定行和列的显示标签
disp = ConfusionMatrixDisplay(
    confusion_matrix=cm,
    display_labels=[["M", "F", "M?", "F?"], ["M", "F"]]
)
disp.plot()
plt.show()

关键说明 #

• 核心矛盾是真实类别有4种，但预测仅输出2种，sklearn默认逻辑会对齐所有出现过的类别，因此必须手动限定矩阵的行、列维度对应的类别范围。
• 第一种方法更贴合二分类器的本质逻辑，把非优质数据归到对应二分类类别；第二种方法保留了真实标签的4个维度，能直观展示不同真实类别被预测为两类的分布情况。

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