Mnist: get confusion matrix

Question

I try to get the confusion matrix of mnist dataset.

It's my code:

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)

The problem is that test_prediction is a 10000 x 10 matrix, while y_test is 10000 x 1 matrix. In fact, the neural net doesn't provide an output for every test sample. How can I calculate confusion matrix for such cases?

And how can I then present the confusion matrix? Can I use the sci-kit library for this purpose?

Answer 1

This is probably because your prediction consist of probabilities of all the possible classes. You need to pick the class with highest probability, which will result in the same dimension as y_test. You can use the argmax() method from numpy. It works something like this:

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])

You can use sklearn for generating confusion matrix. Your code will become something like this

from sklearn.metrics import confusion_matrix
import numpy as np

confusion = confusion_matrix(y_test, np.argmax(test_predictions,axis=1))