Make the learning curve of an ML model more informative

1 minute read

In this short geeky coffee break, we will see how to report the classification loss per epoch for each class separately for a neural network built in Keras.


In multiclass classification scenarios, when training a machine learning model the loss after every epoch is usually an average over all inputs. Instead, we can easily change this and report the loss for every class separately.


The averaged loss hides the performance of the classifier and the learning curve it has for different classes. We have no clue what classes the model has learnt early or later in the training process. Having this information helps us in understanding the learning curve.


In Keras this can be easily achieved using the Callbacks API. The callback simply splits the data into buckets corresponding to the different classes. Then, it calculates the loss for each class separately.

For instance, if we build a simple neural network built in Keras and train it on the CIFAR-10 dataset we get a plot of the loss per class over the epochs as shown below:

Plot of all the losses per class per epoch
Figure 1: Loss per class per epoch

We can observe that the model is learning class 1 (which represents an automobile) quicker than class 3 (which represents a cat).

The source code for the implementation can be found on GitHub. If this is something you like and would like to see similar content you could follow me on LinkedIn or Twitter. Additionally, you can subscribe to the mailing list below to get similar updates from time to time.

Leave a comment