BrianaAI

안녕하세요.

브리아나입니다.

오늘은 yolov8 모델 freeze 하는 방법을 알아보려고 하는데요!

실패

yolo mode=train model=model.pt data=product.yaml imgsz=640 epochs=200 batch=16 freeze=10

이렇게 하니까

SyntaxError: 'freeze' is not a valid YOLO argument.

에러가 뜨더라구요..

그래서 열심히 구글링 해보던 중 다음과 같은 해결 방법의 링크를 발견했는데요!

https://github.com/ultralytics/ultralytics/issues/793

역시 욜로 커뮤니티!

from ultralytics import YOLO

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.")
model=YOLO("model.pt")
model.add_callback("on_train_start", freeze_layer)
results = model.train(data="product.yaml", epochs=200, imgsz=640, batch=8)

main.py 파일을 하나 만들어서 위 코드를 복붙하고

실행시키면

freeze 잘 된 것 같습니다!!

그럼 행복한 코딩 하세요~~~~