DEV Community

Discussion on: Understanding the Confusion Matrix

Collapse
 
overrideveloper profile image
Banso D. Wisdom • Edited

Hi, Manza. I've not worked with Pytorch but I believe it should be similar to the example I gave.

From the snippet you gave, since you already have the confusion matrix (confusion_matrix), what you need to do is create a dictionary of all the classes like this:

dict = { 0: 'Class 1', 1: 'Class 2', 2: 'Class 3',..., n: 'Class n' }

Then pass your confusion matrix (confusion_matrix) and the dictionary of classes (dict) to the function plot_confusion_matrix in the example I gave or a similar function, probably one you've written to your preferences.

That should output a visualization of your confusion matrix.