Fine tune a pre-trained model to add new custom classes

Hello, I have very less knowledge about fine tuning a Tensorflow Model, I want to build an object detection model which can have both the pre-trained classes and new custom classes but I can not find anything to understand the same, it would be great if someone could provide some resources or something. Thank you in advance.


Welcome to the TF Forum!

You can do it in two ways depending on how you use the pre trained weights:

  1. Transfer learning - This technique will only trains the detection head and not the backbone of the network . This means that the backbone weights are sherd between your model and the original model. The idea is that because the backbone is the same for both we can use the backbone to extract the features for the image and then feed each detection head with the features.

If you use this technique then you could do the following:

  • Train the network on your data and save the new detection head weights.
  • Create a new network that has the same backbone but two detection heads, one with the original weights and the second head with the new weights.
  1. A starting point for fitting the entire model - In this technique you train all the network, backbone + detection head. This means that you have two different models.
    This way you’ll load both networks and run inference on the image once with the original network and second with your new network. Then you combine your results.

Let us know if this solves your query!