3
$\begingroup$

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?

$\endgroup$
0
1
$\begingroup$

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.)

  1. 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.

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.