Skip to main content
 首页 » 编程设计

python之MNIST : get confusion matrix

2025年05月04日14telwanggs

我尝试获取 mnist 数据集的混淆矩阵。

这是我的代码:

mnist = tf.keras.datasets.mnist 
 
(x_train, y_train),(x_test, y_test) = mnist.load_data() 
x_train, x_test = x_train / 255.0, x_test / 255.0 
 
model = tf.keras.models.Sequential([ 
  tf.keras.layers.Flatten(), 
  tf.keras.layers.Dense(512, activation=tf.nn.tanh), 
  tf.keras.layers.Dropout(0.2), 
  tf.keras.layers.Dense(10, activation=tf.nn.softmax) 
]) 
model.compile(optimizer='adam', 
              loss='sparse_categorical_crossentropy', 
              metrics=['accuracy']) 
 
 
model.fit(x_train, y_train, epochs=1, callbacks=[history]) 
 
test_predictions = model.predict(x_test) 
 
 
# Compute confusion matrix 
confusion = tf.confusion_matrix(y_test, test_predictions) 

问题是 test_prediction 是一个 10000 x 10 矩阵,而 y_test 是 10000 x 1 矩阵。事实上,神经网络不会为每个测试样本提供输出。对于这种情况,我该如何计算混淆矩阵?

然后我该如何呈现混淆矩阵?我可以为此目的使用 sci-kit 库吗?

请您参考如下方法:

这可能是因为您的预测包含所有可能类别的概率。您需要选择概率最高的类别,这将导致与 y_test 具有相同的维度。您可以使用 numpy 中的 argmax() 方法。它的工作原理是这样的:

import numpy as np 
a = np.array([[0.9,0.1,0],[0.2,0.3,0.5],[0.4,0.6,0]]) 
np.argmax(a, axis=0) 
array([0, 2, 1]) 

您可以使用 sklearn 生成混淆矩阵。你的代码会变成这样

from sklearn.metrics import confusion_matrix 
import numpy as np 
 
confusion = confusion_matrix(y_test, np.argmax(test_predictions,axis=1))