What is the best way to extend an existing model with additional classes of objects?

[AlexPasqua]

Hello,
I have trained a YOLO (v8) model with my custom data, let's say it can recognize 2 types of objects. Now I want to have a model that not only recognizes those 2, but also a third type. I have some pictures of the new object.

What would you say is the best way to obtain that? Is there a way I can train the existing model ONLY on the new data (only new object)? Maybe with the help of some weight freezing, regularization, etc?

Thanks for advice :)

[UltralyticsAssistant]

👋 Hello @AlexPasqua, thank you for your interest in Ultralytics 🚀! Your question about extending an existing model is a great one.

If this is a custom training ❓ question, please provide as much information as possible, such as examples of your dataset (including old and new object examples), your training logs, and any specifics on your current training settings. This will help us understand your current setup better. Ensuring you're following best practices for training, such as balanced datasets and augmentation, is also recommended.

If you're encountering any issues or unexpected results during this process, and if it's a 🐛 bug, please share a minimum reproducible example (MRE) for us to properly investigate.

To ensure you're operating at peak performance, please confirm you're using the latest version of the ultralytics package installed in a Python environment that meets all the necessary requirements, such as Python>=3.8 and PyTorch>=1.8. You can upgrade your setup with the following command:

pip install -U ultralytics

Additionally, YOLO can be run in diverse, verified environments with dependencies pre-installed, like notebook platforms (e.g., Colab, Kaggle), or setups using cloud services (e.g., GCP, AWS) or Docker Images.

This is an automated response, but no worries, an Ultralytics engineer will assist you shortly 😊!

[glenn-jocher]

@AlexPasqua thank you for your question! To extend your YOLOv8 model with a new class, the best approach is to fine-tune it using a combined dataset of your original classes and the new class. This ensures the model retains knowledge of the original classes while learning the new one. Direct training on only the new class is not recommended as it may cause the model to forget previously learned classes (catastrophic forgetting). You can use the train mode with your combined dataset and adjust parameters like freeze if you wish to freeze specific layers. For more details, refer to the Train documentation.

[Y-T-G]

If you want to keep existing performance, you will have to add new parameters and train only those while freezing the old ones. I do it here by adding a new head. The caveat is that this reduces the model's speed.

https://y-t-g.github.io/tutorials/yolov8n-add-classes/

[AlexPasqua]

If you want to keep existing performance, you will have to add new parameters and train only those while freezing the old ones. I do it here by adding a new head. The caveat is that this reduces the model's speed.

https://y-t-g.github.io/tutorials/yolov8n-add-classes/

Hi, thanks for the advice!

What about freezing all the weights, then scraping the existing head entirely and replacing it with a new one that then gets trained with the new data plus the old data (only some of the old data to keep the 3 classes balanced) ?
What do you think about that? Would that be simpler than what you suggested?

[glenn-jocher]

@AlexPasqua your approach of freezing the backbone weights and replacing the head to train on a balanced dataset with old and new classes can work well. It simplifies the process and helps retain previous performance while learning the new class. Just ensure the dataset is well-balanced and large enough to avoid overfitting.

[Y-T-G]

@AlexPasqua There's no guarantee you would retain the original performance on the old classes.

[AlexPasqua]

@AlexPasqua There's no guarantee you would retain the original performance on the old classes.

I agree, but maybe there is a good chance to obtain good enough performance.

Anyway, I'll experiment a little bit. I read your guide on how to add an extra head, it's very useful. Maybe I'll ask some questions here in case I have any (probably)

[Y-T-G]

@AlexPasqua You can ask in the comments in the guide.

You can also check this.

#18710