How to freeze the backbone, feature extraction layers, in YOLO v8

Question

I want to train the YOLO v8 in transfer learning on my custom dataset.
I have different classes than the base training on the COCO dataset.

Yet I don't want to learn again the feature extraction. Hence I though following the Ultralytics YOLOv8 Docs - Train.

Yet, When I train on my small dataset I want to freeze the backbone.

How can I do that?

I looked at the documentation and couldn't find how to do so.

Answer 1

You can do the following

def freeze_layer(trainer):
    model = trainer.model
    num_freeze = 10
    print(f"Freezing {num_freeze} layers")
    freeze = [f'model.{x}.' for x in range(num_freeze)]  # layers to freeze 
    for k, v in model.named_parameters(): 
        v.requires_grad = True  # train all layers 
        if any(x in k for x in freeze): 
            print(f'freezing {k}') 
            v.requires_grad = False 
    print(f"{num_freeze} layers are freezed.")

Then add this function as a custom callback function to the model

model = YOLO("yolov8x.pt") 
model.add_callback("on_train_start", freeze_layer)
model.train(data="./dataset.yaml")

Original answer is provided in one of the issues in ultralytics repo Freezing layers yolov8 #793