sgrvinod / a-PyTorch-Tutorial-to-Object-Detection

SSD: Single Shot MultiBox Detector | a PyTorch Tutorial to Object Detection
MIT License
3.04k stars 718 forks source link

Fine-tune to specific class? #64

Closed Awwh closed 4 years ago

Awwh commented 4 years ago

This code is working well for me with the classes from the VOC dataset but I would like to fine-tune this model so that it only looks for people.

Is there any easy way to do this?

zackdilan commented 4 years ago

how good is your current model in detecting people (I assume people are one of the classes in VOC dataset)? if the current predictions are bad, you can fine-tune the current model with minimal changes in the model definition (where you have to only deal with two classes people or background), and of course, you should make sure you adapt the other ancillary functions too.

Awwh commented 4 years ago

It is relatively good as people are included in the dataset as you assumed, however I do have my own smaller dataset with more accurate use case images I would like to fine-tune the model to.

You mentioned altering the model file to only deal with the 2 classes which makes sense however I cant seem to locate the piece of code that needs to be changed. Can you assist me with this?

Thanks for the help so far!

zackdilan commented 4 years ago

okay! I am afraid that I haven't gone through the code really far, But from a glimpse, I have understood, you don't want to make any changes in the model.py, in the sense, you just need to change the number of classes (variable ) passed into the model definition.

Then obviously adapt your data loader accordingly, I hope it's not difficult as it seems as the author has written the code in a well-generalized manner.

sgrvinod commented 4 years ago

Hey @Awwh, this should be easy, I'll take a look at the code and let you know. I think you mailed me too, but I was really busy and couldn't reply. Sorry for the delay. I'll check and let you know in 24 hours. Also, thank you @zackdilan!

Awwh commented 4 years ago

Thanks for all the help guys

sgrvinod commented 4 years ago

Hi @Awwh,

If you need to train it from scratch on the VOC dataset, but only on person:

  1. Edit this line to include only one class, i.e. {'person'}. This means that the label map will ultimately contain only two classes background and person. The index of the the person class in the label map is 1 (background is 0}.

  2. Make sure that this line is correct in your version of the code, because I just updated it.

I don't think you need to do anything else, because this line filters out objects that are not present in the label_map, and this line filters out images that don't contain any objects that are present in the label_map. This line automatically accounts for the new target task (i.e. a 2 class problem, background vs. person) as defined in the label_map.

And then you can begin training from scratch. After you train, you can evaluate your model and check if the Average Precision (AP) for the person class in your new model is better than what it used to be when trained on all classes.

If you need to train it from scratch on your new dataset, which contains only class x (person):

Here, class x is whatever people are called in your new dataset.

  1. The same procedure as above, except edit this line to include class x instead of person.

  2. You'd need to edit this function to work on your new dataset and create the image and object JSON files for your new dataset.

If you need to first train on the VOC task and then fine-tune to class x (person) in your new dataset:

  1. After you train on the VOC task (all classes or just the person class), edit your new dataset such that class x is now called person in this new dataset.

  2. Then you'd need to edit this function to work on your new dataset and create the image and object JSON files for your new dataset.

  3. And then begin training once again on the new data, starting with the trained VOC model checkpoint, and with a small learning rate. This should fine-tune your model to the new dataset.

You could try to first train from scratch on only the person class in the VOC dataset, and the fine-tune to your new dataset for better performance.

If you want to initially train on all VOC classes, but you don't want it to retain the prediction layers that predict non-person classes in the final fine-tuned version, you can try removing the prediction layers from the VOC model checkpoint and adding new prediction layers for a 2-class problem before fine-tuning on your new dataset, but this is more complicated and more than I can address in this post.

Hope this helps.

Awwh commented 4 years ago

Thank you so much for this @sgrvinod! It's honestly gonna be so helpful.

I haven't had a chance to apply this myself yet but I'm gonna close this now as it looks like its very well solved.

Thankyou