如何让sklearn.metrics.confusion_matrix()始终返回TP,TN,FP,FN?

时间:2017-09-15 00:25:57

标签: python scikit-learn confusion-matrix

我正在使用 sklearn.metrics.confusion_matrix(y_actual, y_predict) 来提取tn,fp,fn,tp,并且大部分时间它都能正常运行。

from sklearn.metrics import confusion_matrix

y_actual, y_predict = [1,1,1,1], [0,0,0,0]
tn, fp, fn, tp = confusion_matrix(y_actual, y_predict).ravel()
>>> [0 0 4 0]   # ok

y_actual, y_predict = [1,1,1,1],[0,1,0,1]
tn, fp, fn, tp = confusion_matrix(y_actual, y_predict).ravel()
>>> [0 0 2 2]   # ok

但是,在某些情况下,confusion_matrix()并不总是返回这些信息,我会得到ValueError,如下所示。

from sklearn.metrics import confusion_matrix

y_actual, y_predict = [0,0,0,0],[0,0,0,0]
tn, fp, fn, tp = confusion_matrix(y_actual, y_predict).ravel()
>>> [4]    # ValueError: not enough values to unpack (expected 4, got 1)

y_actual, y_predict = [1,1,1,1],[1,1,1,1]
tn, fp, fn, tp = confusion_matrix(y_actual, y_predict).ravel()
>>> [4]    # ValueError: not enough values to unpack (expected 4, got 1)

我的临时解决方案是编写自己的函数来提取这些信息。有什么方法可以强制 confusion_matrix() 始终返回tn,fp,fn,tp输出?

由于

1 个答案:

答案 0 :(得分:7)

此问题与输入矩阵中包含的唯一标签数量有关。在你的第二个例子中,它(正确地)构建一个只有一个类的混淆矩阵,分别为0或1。

要强制它输出两个类,即使没有预测其中一个类,也要使用label属性。

y_actual, y_predict = [0,0,0,0],[0,0,0,0]
tn, fp, fn, tp = confusion_matrix(y_actual, y_predict, labels=[0,1]).ravel()
>> array([[4, 0],
          [0, 0]])