Model fine-tuning - Train a specific class of the pretrained YOLOv8 model without compromising information related to other classes

[Masoudvahid]

I am using the pretrained yolov8 model trained on coco dataset, but there is a case that the model does not perform ideally on a specific class.
So I created the desired dataset for this specific class. Then I started to train on the new dataset (300 images) and by freezing the first 10 layers (as far as I understood first 10 layers are backbone).
The result of training is acceptable, but now I am missing the all other 78 classes.
My question is how can I fine-tune the model without loosing other classes and only training on very small number of images (without training on coco dataset again).

The command which I used is as follows:
yolo detect train imgsz=640 batch=16 epochs=50 data="dataset.yaml" model=yolov8x.pt freeze=10

where dataset.yaml is the path to the new dataset and all default coco classes.

It worthsnoting that when I check the number of classes on the output later of the model, I see all classes, but the model predicts none of them.

And lastly, I also freezed all 22 layers, and retrained. Just in case of making sure if any of the classes would resist, but all the classes were predicting nothing.

[pderrenger]

@Masoudvahid when fine-tuning a pretrained model on a new dataset for a specific class, it's important to ensure that the new dataset includes examples of all the classes you want the model to recognize. If you train only on the new class, the model will likely forget the other classes, a phenomenon known as catastrophic forgetting.

To fine-tune on a specific class without losing performance on the other classes, you should include examples of the other classes in your training dataset. This can be a mix of new images for the class of interest and a subset of the original dataset images for the other classes. This way, the model maintains its ability to recognize the other classes while improving on the new class.

Here's a general approach to fine-tuning your model:

1. Create a Mixed Dataset: Combine your new images for the specific class with a representative subset of the original COCO dataset that includes all classes.

2. Use a Lower Learning Rate: When fine-tuning, use a lower learning rate to make smaller updates to the weights, which helps in retaining the previously learned information.

3. Selective Freezing: Instead of freezing the first few layers, you might consider freezing only the layers that are less likely to need retraining, such as the early convolutional layers. The later layers, which are more class-specific, can be fine-tuned.

4. Class Weights: Adjust the loss function to give more weight to the new class during training, which can help the model learn from the smaller number of examples without forgetting the other classes.

5. Regularization Techniques: Use techniques like dropout or data augmentation to prevent overfitting to the new class.

6. Monitor Validation Performance: Ensure you have a validation set that includes all classes. Monitor the performance on this set to ensure that the model retains its ability to detect all classes.

Here's an example command that incorporates some of these principles (note that you'll need to adjust the learning rate and other hyperparameters according to your specific needs):

from ultralytics import YOLO

# Load the pretrained model
model = YOLO('yolov8x.pt')

# Fine-tune the model on the mixed dataset with a lower learning rate
results = model.train(data='mixed_dataset.yaml', epochs=50, lr0=0.001, freeze=10)

Remember to create a mixed_dataset.yaml that includes paths to both the new class images and a subset of the original dataset. Adjust the learning rate (lr0) as needed. The freeze parameter should be set based on the layers you want to freeze, if any.

Keep in mind that fine-tuning is a delicate process, and it may take some experimentation to get the balance right between retaining old knowledge and learning new information. 🧪🔍

[xihri412]

@pderrenger For the mixed dataset, should I label the old classes in the new dataset and label the new classes in the COCO dataset?

[pderrenger]

Hi @xihri412! Yes, for effective fine-tuning, you should label the old classes in your new dataset. This helps maintain the model's ability to recognize these classes. There's no need to label new classes in the original COCO dataset unless those new classes are also present and you want to improve detection for them as well. Essentially, ensure that your training data is representative of all classes you want the model to detect. Happy fine-tuning! 😊