Plotting confusion matrix for multi classification problem

Question

I am using google colab to solve a multi-classification problem. I am trying to plot the confusion matrix for this problem, I have tried doing so using :

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

confusion_matrix=multilabel_confusion_matrix(test_labels, y_pred)
disp=ConfusionMatrixDisplay(confusion_matrix,display_labels=test_labels)
disp.plot(include_values=True, cmap="viridis", ax=None, xticks_rotation="vertical")
plt.show()

However, I am getting this error :

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
<ipython-input-25-3fe075d4503c> in <module>()
      5 confusion_matrix=multilabel_confusion_matrix(test_labels, y_pred)
      6 disp=ConfusionMatrixDisplay(confusion_matrix,display_labels=test_labels)
----> 7 disp.plot(include_values=True, cmap="viridis", ax=None, xticks_rotation="vertical")
      8 plt.show()

5 frames
/usr/local/lib/python3.7/dist-packages/matplotlib/image.py in set_data(self, A)
    697                 or self._A.ndim == 3 and self._A.shape[-1] in [3, 4]):
    698             raise TypeError("Invalid shape {} for image data"
--> 699                             .format(self._A.shape))
    700 
    701         if self._A.ndim == 3:

TypeError: Invalid shape (8, 2, 2) for image data

Answer 1

The result of multilabel_confusion_matrix is many matrices. ConfusionMatrixDisplay only displays a single matrix. One option is to loop through the matrices to plot each one separately. Something like this:

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


test_labels = [0, 1, 2, 1, 0]
y_pred      = [0, 1, 2, 1, 1]
confusion_matrices = multilabel_confusion_matrix(test_labels, y_pred)
for confusion_matrix in confusion_matrices:
    disp = ConfusionMatrixDisplay(confusion_matrix, display_labels=test_labels)
    disp.plot(include_values=True, cmap="viridis", ax=None, xticks_rotation="vertical")
    plt.show()