How to add a new category to a deep learning model?

Question

Say I have done transfer learning on a pre-trained network to recognize 10 objects. How can I add a $11^{th}$ item that the network can classify without losing all the 10 categories I already trained and the information from the original pre-trained model? A friend told me that active research is going on in this field, but I cannot find any relevant papers or a names to search for?

Answer 1

If this is just a one-time case, you can simply re-train the neural network. If you frequently have to add new classes, then this is a bad idea. What you want to do in such cases is called content-based image retrieval (CBIR), or simply image retrieval or visual search. I will explain both cases in my answer below.

One-time case

If this just happens once - you forgot the 11th class, or your customer changed his/her mind - but it won't happen again, then then you can simply an 11th output node to the last layer. Initialize the weights to this node randomly, but use the weights you already have for the other outputs. Then, just train it as usual. It might be helpful to fix some weights, i.e. don't train these.

An extreme case would be to only train the new weights, and leave all others fixed. But I am not sure whether this will work that well - might be worth a try.

Content-based image retrieval

Consider the following example: you are working for a CD store, who wants their customers to be able to take a picture of an album cover, and the application shows them the CD they scanned in their online store. In that case, you would have to re-train the network for every new CD they have in the store. That might be 5 new CDs each day, so re-training the network that way is not suitable.

The solution is to train a network, which maps the image into a feature space. Each image will be represented by a descriptor, which is e.g. a 256-dimensional vector. You can "classify" an image by calculating this descriptor, and comparing it to your database of descriptors (i.e. the descriptors of all CDs you have in your store). The closest descriptor in the database wins.

How do you train a neural network to learn such a descriptor vector? That is an active field of research. You can find recent work by searching for keywords like "image retrieval" or "metric learning".

Right now, people usually take a pre-trained network, e.g. VGG-16, cut off the FC layers, and use the final convolutional as your descriptor vector. You can further train this network e.g. by using a siamese network with triplet loss.

Answer 2

This can be done easily.

First build a model with those 10 classes and save the model as base_model.

Load the base_model and also define a new model named new_model as-

new_model = Sequential()

Then add the layers of the base_model to the new_model -

# getting all the layers except the last two layers
for layer in base_model.layers[:-2]: #just exclude the last two layers from base_model
    new_model.add(layer)

Now make the layers of the new model non-trainable as you don't want your model to be trained again.

# prevent the already trained layers from being trained again
for layer in new_model.layers:
    layer.trainable = False

Now as you transfer learning, when you remove the last layers, the model kind of forgets about the 10 classes so we have to retain the weights of the base_model to the new_model -

weights_training = base_model.layers[-2].get_weights()
new_model.layers[-2].set_weights(weights_training) 

Now add a dense layer at the end and we will only train this dense layer in this example.

new_model.add(Dense(CLASSES, name = 'new_Dense', activation = 'softmax'))

Now train the model and I hope it gives the right output for all the 11 classes.

Happy Learning.

Answer 3

Your network topology might look different, but in the very end, your pre-trained network has a layer, which handles the recognition of 10 original classes. The easiest (and working) trick to introduce the 11th, 12th.. nth class, is to use all the layers before the last as granted and add an additional layer (in a new model, or as a parallel one) that will also sit on top of all but last layers, will be looking alike to the 10class layer (which is most probably matmul of dense layer and a matrix of shape [len(dense layer), 10] with optional bias).

Your new layer would be a matmul layer with shape [len(dense layer), len(new classes)].

Without access to original training data, you would have two options:

1. Freeze all the weights in original layers by allowing "new" model to optimize only new weights. That will give you exactly same predictive power for original 10 classes and might give OK performance for new ones.
2. Train whole network at once (by propagating error of new classes), which might be working for new class(es), but you will end up with ineffective original solution for 10 classes (since weights will be changed for the lower classes and final layer won't be updated to match those changes).

Although, given you have access to original training data, you can easily add new class to the original network and re-train it to support 11 classes out of the box.