Incremental Classifier and Representation Learning in Yolo models

Question

I'm encountering an issue with my YOLO model.

Initially, I trained it with 7 classes. Now, I want to add 4 new classes to the model. However, when I combine the data for the original 7 classes with the new 4 classes, the training time and associated cloud costs significantly increase. What's a good solution to efficiently incorporate these additional classes into the model without inflating training time and costs?

My expecting is reduce the cost and training time in incremental learnng.

Answer 1

For adding new classes to an already trained model, you might consider the concept of Transfer Learning. Instead of retraining the entire model from scratch with both old and new classes combined, you can:

1. Freeze the layers up to the last one or few (which have learned feature representations from your initial training).
2. Then, only train the final layers (or add new ones) to learn the additional classes.

Here's a simplified code snippet to give you an idea:

from ultralytics import YOLO

model = YOLO('runs/detect/train/weights/best.pt') # load a pretrained model (recommended for training)
 
# unfreeze the last 4 layers of model,
named_parameters = list(model.named_parameters())
for i, (name, param) in enumerate(named_parameters):
    if i >= len(named_parameters) - 4:  # Unfreeze the last 4 layers
        param.requires_grad = True
    else:
        param.requires_grad = False


# # Proceed with training on the new dataset (with 11 classes now)
model.train(data="dataset/data.yaml", epochs=200, batch=16, workers=1)  # train the model

This code illustrates the general approach. You'll need to adjust the specifics based on your actual model structure and how your data is organized.

Remember to update your dataset configuration file (new_data_with_11_classes.yaml in the example) to reflect all 11 classes.

This approach can significantly reduce both training time and costs.

@glenn-jocher