Custom metrics in  tensorflow2.2  | Towards Data Science

Keras has simplified DNN based machine learning a lot and it keeps getting better. Here we show how to implement metric based on the confusion matrix (recall, precision and f1) and show how using them is very simple in tensorflow 2.2. You can directly run the notebook in Google Colab.
When considering a multi-class problem it is often said that accuracy is not a good metric if the classes are imbalanced. While that is certainly true, accuracy is also a bad metric when all classes do not train equally well even if the datasets are balanced.
In this article, I will use Fashion MNIST to highlight this aspect. This is however not the only goal of this article as this can be done by simply plotting the confusion matrix on the validation set at the end of training. What we discuss here is the ability to easily extend keras.metrics.Metric class to make a metric that tracks the confusion matrix during training and can be used to follow the class specific recall, precision and f1 and plot them in the usual way with keras.
Getting class specific recall, precision and f1 during training is useful for at least two things:
We can see if training is stable — plots do not jump around too much — on a per class basisWe can implement more customized training based on class statistic — early stopping or even dynamically changing class weights.
We can see if training is stable — plots do not jump around too much — on a per class basis
We can implement more customized training based on class statistic — early stopping or even dynamically changing class weights.
Furthermore, since tensorflow 2.2, integrating such custom metrics into training and validation has become very easy thanks to the new model methods train_step and test_step. There is also an associate predict_step that we do not use here but works in the same spirit.
So lets get down to it. We first make a custom metric class. While there are more steps to this and they are show in the referenced jupyter notebook, the important thing is to implement the API that integrates with the rest of Keras training and testing workflow. That is as simple as implementing and update_state that takes in the true labels and predictions, a reset_states that re-initializes the metric.
class ConfusionMatrixMetric(tf.keras.metrics.Metric):                def update_state(self, y_true, y_pred,sample_weight=None):        self.total_cm.assign_add(self.confusion_matrix(y_true,y_pred))        return self.total_cm            def result(self):        return self.process_confusion_matrix()        def confusion_matrix(self,y_true, y_pred):        """        Make a confusion matrix        """        y_pred=tf.argmax(y_pred,1)        cm=tf.math.confusion_matrix(y_true,y_pred,dtype=tf.float32,num_classes=self.num_classes)        return cm        def process_confusion_matrix(self):        "returns precision, recall and f1 along with overall accuracy"        cm=self.total_cm        diag_part=tf.linalg.diag_part(cm)        precision=diag_part/(tf.reduce_sum(cm,0)+tf.constant(1e-15))        recall=diag_part/(tf.reduce_sum(cm,1)+tf.constant(1e-15))        f1=2*precision*recall/(precision+recall+tf.constant(1e-15))        return precision,recall,f1      
In the normal Keras workflow, the method result will be called and it will return a number and nothing else needs to be done. However, in our case we have three tensors for precision, recall and f1 being returned and Keras does not know how to handle this out of the box. This is where the new features of tensorflow 2.2 come in.
Since tensorflow 2.2 it is possible to modify what happens in each train step (i.e. training on a mini-batch) transparently (whereas earlier one had to write an unbounded function that was called in a custom training loop and one had to take care of decorating it with tf.function to enable autographing).
Again, details are in the referenced jupyter notebook but the crux is the following
class MySequential(keras.Sequential):        def train_step(self, data):        # Unpack the data. Its structure depends on your model and        # on what you pass to `fit()`.        x, y = data        with tf.GradientTape() as tape:            y_pred = self(x, training=True)  # Forward pass            # Compute the loss value.            # The loss function is configured in `compile()`.            loss = self.compiled_loss(                y,                y_pred,                regularization_losses=self.losses,            )        # Compute gradients        trainable_vars = self.trainable_variables        gradients = tape.gradient(loss, trainable_vars)        # Update weights        self.optimizer.apply_gradients(zip(gradients, trainable_vars))        self.compiled_metrics.update_state(y, y_pred)        output={m.name: m.result() for m in self.metrics[:-1]}        if 'confusion_matrix_metric' in self.metrics_names:            self.metrics[-1].fill_output(output)        return output
That’s it. Now you can create (using the above class not keras.Sequential), compile and fit a sequential model (the procedure to do with with Functional and Subclassing API is straightfoward and one just implements the above function). The resulting history now has elements like val_F1_1 etc.
The advantage of this is that we can see how individual classes train
We see that class 6 trains pretty bad with an F1 of around .6 on the validation set but the training itself is stable (the plot doesn’t jump around too much).
As mentioned in the beginning, getting the per-class metrics during training is useful for at least two things:
We can see if training is stable — plots do not jump around too much — on a per class basisWe can implement more customized training based on class statistic based early stopping or even dynamically changing class weights.
We can see if training is stable — plots do not jump around too much — on a per class basis
We can implement more customized training based on class statistic based early stopping or even dynamically changing class weights.
Finally, let's look at the confusion matrix to see what is happening with class 6
In the confusion matrix, true classes are on the y-axis and predicted ones on the x-axis. We see that shirts (6), are being incorrectly labeled mostly as t-shirts (0), pullovers(2) and coats (4). Conversely, the mislabelling as shirts happens mostly for t-shirts. These kinds of mistakes are reasonable and I will discuss in a separate article what can be done to improve training in such cases.
Ref:
Code on githubGoogle Colab notebook
Code on github
Google Colab notebook