When doing transfer learning, which initial layers do we need to freeze, and how should I change the last layer for my task?

Question

I want to train a neural network for the detection of a single class, but I will be extending it to detect more classes. To solve this task, I selected the PyTorch framework.

I came across transfer learning, where we fine-tune a pre-trained neural network with new data. There's a nice PyTorch tutorial explaining transfer learning. We have a PyTorch implementation of the Single Shot Detector (SSD) as well. See also Single Shot MultiBox Detector with Pytorch — Part 1.

This is my current situation

• The data I want to fine-tune the neural network with is different from the data that was used to initially train the neural network; more specifically, the neural network was initially trained with a dataset of 20 classes

• I currently have a very small labeled training dataset.

To solve this problem using transfer learning, the solution is to freeze the weights of the initial layers, and then train the neural network with these layers frozen.

However, I am confused about what the initial layers are and how to change the last layers of the neural network to solve my specific task. So, here are my questions.

1. What are the initial layers in this case? How exactly can I freeze them?

2. What are the changes I need to make while training the NN to classify one or more new classes?

Answer 1

After a lot of browsing online for answers to these questions, this is what I came up with.

1. What are the initial layers in this case? How exactly can I freeze them?

The initial few layers are said to extract the most general features of any kind of image, like edges or corners of objects. So, I guess it actually would depend on the kind of backbone architecture you are selecting.

How to freeze the layers depends on the framework we use.

(I have selected PyTorch as the framework. I found this tutorial Some important Pytorch tasks - A concise summary from a vision researcher, which seems to be useful.)

2. What are the changes I need to make while training the NN to classify one or more new classes?

I just need to change just the number of neurons (or units) in the final layer to the number of classes/objects your new dataset contains.