YOLOv8冻结层的迁移学习 全网首发

虚假的冻结层

from ultralytics import YOLO

model = YOLO('yolov8x.pt')
data = '/home/tm/fyjx/FangTeDatasets/Datasets/data.yaml'

# for k, v in model.model.named_parameters():
#     print(k)
freeze = 10
freeze = [f'model.{x}.' for x in range(freeze)]  # layers to freeze
for k, v in model.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

model.train(data = data, epochs = 50, batch = 8, save = True, device = 4)

model.train在调用时会创建trainer类,请看vcr:

self.trainer.model = self.trainer.get_model(weights=self.model if self.ckpt else None, cfg=self.model.yaml)

此时的模型配置更新了,冻结了个p

真正的冻结层

YOLO 提供了丰富的回调函数, 我们在创建trainer后冻结层, 在训练程序开始时调用函数冻结真正的model, 那就是trainer.model

from ultralytics import YOLO

model = YOLO('yolov8x.pt')
data = '/app/Datasets/data.yaml'

def freeze_model(trainer):
    # Retrieve the batch data
    model = trainer.model
    print('Befor Freeze')
    for k, v in model.named_parameters():
        print('\t', k,'\t', v.requires_grad)
        
        
    freeze = 10
    freeze = [f'model.{x}.' for x in range(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('After Freeze')
    for k, v in model.named_parameters():
        print('\t', k,'\t', v.requires_grad)
        

model.add_callback("on_pretrain_routine_start", freeze_model)

model.train(data = data, epochs = 50, batch = 32, save = True)