Customize Classification Model Output Layer

Save classification labels and top confidences in a custom layer using Keras

Image classification is a handbook example of deep learning applications. The standard way of making a classification model involves a preprocessing step where the human-readable class labels (eg.: ‘car’, ‘person’, ‘cat’) are changed to machine-ready numbers (eg.: 0,1,2). The most common way of doing this is to map the list of possible classes with their indices. Of course, this requires a postprocessing step as well, where the result is converted to the expected form. A common way is to store the label of the class with the highest score and the score (a widely used term for this is confidence ).

In this story, I will show an example of storing the labels in the model using a custom layer at the end of the model. The preferred output of my model is a list of the top k labels and their confidence scores. The formated output can be useful in production where the product is a model, and the labels have to be stored inside the model. Or, when the list of the labels changes with each iteration of models.

Input: image(W,H,C)
Outputs: labels(k) string, confidences(k) float

Training a Classification Model

For this story, I will use a simple classification model. This Colab notebook shows an example of a classifier trained on the Fashion MNIST dataset (trained on 60000 images and tested on 10000). The model expects 28x28x1 grayscale images and returns with a softmax probability of 10 classes. The list of the class labels are:

class_names = ['T-shirt/top', 'Trouser', 'Pullover', 'Dress', 'Coat', 'Sandal', 'Shirt', 'Sneaker', 'Bag', 'Ankle boot']

A straightforward way to handle the output of the model can be a simple mapping: finding the index with the most score and use the label of that index.

class_names[np.argmax(predictions[0])] to get the label of the image 0.

Building a Custom Layer

To understand how to make a custom layer in Keras, I suggest reading the original documentation in Keras and TensorFlow .

I want to implement a custom layer that stores the labels and a topn value, so the output of the layer can be the top n confidence labels and their scores. For this, we have to overwrite the __init__ , the call , compute_output_shape and get_config functions of the layer.

Init

In the init function, we store the constructor parameters as the class’ field. Don’t forget to call the parent class’ constructor! super(LabelLimitLayer, self).__init__(**kwargs)

Call

The call function expects the output of the previous layer as an input and calculates the top_k classes and their labels. For this, I used TensorFlow functions.

To create a (?,len(labels)) shaped tensor from the list of labels, we first create a tensor using the list (parameter of our custom class) and then expand it using the shape of the previous layer’s output (we extract the batch_size from it). The steps are:

tf_labels = tf.constant([self.labels], dtype=”string”) string type tensor

tf_labels = tf.tile(tf_labels,[batch_size,1]) expanding to get (?,1) dynamic shape handling the batches.

To select the top k scores, I used the corresponding TensorFlow function. We store the indices so we can map the labels for those indices as well as the confidence values.

top_k = tf.nn.top_k(x, k=self.topn, sorted=True, name=”top_k”).indices

top_conf = tf.gather(x, top_k, batch_dims=1)
top_labels = tf.gather(tf_labels, top_k, batch_dims=1)

Finally, the layer returns with the last two tensors.

return [top_conf, top_labels]

Compute_output_shape

top_shape = (batch_size, self.topn)
return [top_shape, top_shape]

Get_config

To serialize the custom layer (when saving a model), one has to update the config with the values of the class parameters. Don’t forget to add the superclass’ config to the dictionary!

Adding the Output Layer to the Model

In this example, I added a custom layer at the end of the base classification model, with labels ( class_names ) and top_k value ( 2 ).

label_layer = LabelLimitLayer(class_names, 2)(base_model.output)
label_model = Model(base_model.input, label_layer)

Saving and Loading Model

Finally, to save a model, we can use the Keras model’s save function. To load a model with a custom layer, we have to define this custom layer at the custom_objects parameter.

label_model.save(‘test.h5’)
restored_model = keras.models.load_model(“test.h5”, custom_objects={“LabelLimitLayer”:LabelLimitLayer})

Summary

This snippet shows how to use Keras custom layers to create string labels as the output of a model. The story also uses top_k to keep only the relevant classes. The corresponding codes are available at this Colab notebook .