YOLOv5에서 동결된 레이어를 사용한 전이 학습#

📚 이 가이드에서는 전이 학습을 구현할 때 YOLOv5 🚀 레이어를 동결하는 방법을 설명합니다. 전이 학습은 전체 네트워크를 처음부터 다시 학습하지 않고도 새 데이터로 모델을 빠르게 재학습할 수 있게 해주는 강력한 기계 학습 (ML) 기법입니다. 초기 레이어의 가중치를 동결하고 이후 레이어의 파라미터만 업데이트하면 컴퓨팅 리소스 요구 사항과 학습 시간을 크게 줄일 수 있습니다. 그러나 이 접근 방식은 최종 모델 정확도에 약간 영향을 줄 수 있습니다.

시작하기 전에#

먼저 YOLOv5 리포지토리를 복제하고 requirements.txt에 나열된 필수 종속 항목을 설치합니다. Python>=3.8.0 환경에 PyTorch>=1.8이 설치되어 있는지 확인합니다. 사전 학습된 모델과 필수 데이터셋은 최신 YOLOv5 릴리스에서 자동으로 다운로드됩니다.

git clone https://github.com/ultralytics/yolov5 # clone repository
cd yolov5
pip install -r requirements.txt # install dependencies

레이어 동결 작동 방식#

신경망의 레이어를 동결하면 학습 과정에서 해당 레이어의 파라미터(가중치 및 편향)가 업데이트되지 않도록 할 수 있습니다. PyTorch에서는 레이어 텐서의 requires_grad 속성을 False로 설정하여 이를 구현합니다. 결과적으로 역전파 중에는 이러한 레이어에 대한 그래디언트가 계산되지 않으므로 연산량과 메모리를 절약할 수 있습니다.

YOLOv5는 학습 스크립트에서 다음과 같이 레이어 동결을 구현합니다.

# Freeze specified layers
freeze = [f"model.{x}." for x in range(freeze)]  # Define layers to freeze based on module index
for k, v in model.named_parameters():
    v.requires_grad = True  # Ensure all parameters are initially trainable
    if any(x in k for x in freeze):
        print(f"Freezing layer: {k}")
        v.requires_grad = False  # Disable gradient calculation for frozen layers

모델 아키텍처 살펴보기#

어떤 레이어를 동결할지 결정하려면 YOLOv5 모델의 구조를 이해하는 것이 중요합니다. 다음 Python 코드 조각을 사용하여 모든 모듈의 이름과 해당 파라미터를 확인할 수 있습니다.

# Assuming 'model' is your loaded YOLOv5 model instance
for name, param in model.named_parameters():
    print(name)

"""
Example Output:
model.0.conv.conv.weight
model.0.conv.bn.weight
model.0.conv.bn.bias
model.1.conv.weight
model.1.bn.weight
model.1.bn.bias
model.2.cv1.conv.weight
model.2.cv1.bn.weight
...
"""

YOLOv5 아키텍처는 일반적으로 특징 추출을 담당하는 백본(표준 구성인 YOLOv5s/m/l/x에서는 레이어 0~9)과 객체 탐지를 수행하는 헤드(나머지 레이어)로 구성됩니다.

# Example YOLOv5 v6.0 backbone structure
backbone:
    # [from, number, module, args]
    - [-1, 1, Conv, [64, 6, 2, 2]]  # Layer 0: Initial convolution (P1/2 stride)
    - [-1, 1, Conv, [128, 3, 2]] # Layer 1: Downsampling convolution (P2/4 stride)
    - [-1, 3, C3, [128]]          # Layer 2: C3 module
    - [-1, 1, Conv, [256, 3, 2]] # Layer 3: Downsampling convolution (P3/8 stride)
    - [-1, 6, C3, [256]]          # Layer 4: C3 module
    - [-1, 1, Conv, [512, 3, 2]] # Layer 5: Downsampling convolution (P4/16 stride)
    - [-1, 9, C3, [512]]          # Layer 6: C3 module
    - [-1, 1, Conv, [1024, 3, 2]]# Layer 7: Downsampling convolution (P5/32 stride)
    - [-1, 3, C3, [1024]]         # Layer 8: C3 module
    - [-1, 1, SPPF, [1024, 5]]    # Layer 9: Spatial Pyramid Pooling Fast

# Example YOLOv5 v6.0 head structure
head:
    - [-1, 1, Conv, [512, 1, 1]] # Layer 10
    - [-1, 1, nn.Upsample, [None, 2, "nearest"]] # Layer 11
    - [[-1, 6], 1, Concat, [1]] # Layer 12: Concatenate with backbone P4 (from layer 6)
    - [-1, 3, C3, [512, False]] # Layer 13: C3 module
    # ... subsequent head layers for feature fusion and detection

동결 옵션#

학습 명령에서 --freeze 인수를 사용하여 동결할 레이어를 제어할 수 있습니다. 이 인수는 처음으로 동결되지 않는 모듈의 인덱스를 지정하며, 이 인덱스보다 앞선 모든 모듈의 가중치가 동결됩니다. 특정 블록에 해당하는 인덱스를 확인해야 하는 경우 model.model(nn.Sequential)을 사용하여 모듈 순서를 확인합니다.

백본만 동결#

대규모 데이터셋인 COCO에서 학습한 일반적인 특징 추출 능력을 유지하면서 모델을 새로운 객체 클래스에 적용할 때 흔히 사용하는 방법으로, 전체 백본(레이어 0~9)을 동결하려면 다음과 같이 실행합니다.

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 10

이 전략은 대상 데이터셋이 원래 학습 데이터(예: COCO)와 유사한 저수준 시각적 특징(윤곽선, 텍스처)을 공유하지만 객체 범주는 다른 경우에 효과적입니다.

최종 탐지 레이어를 제외한 모든 레이어 동결#

최종 출력 컨볼루션 레이어(Detect 모듈의 일부이며 일반적으로 마지막 모듈, 예: YOLOv5s의 모듈 24)만 학습 가능하도록 두고 네트워크의 거의 전체를 동결하려면 다음과 같이 실행합니다.

python train.py --weights yolov5m.pt --data your_dataset.yaml --freeze 24

이 접근 방식은 학습된 특징의 대부분을 그대로 유지하면서 출력 클래스 수가 다른 경우에 맞게 모델을 조정해야 할 때 유용합니다. 파인 튜닝에 필요한 컴퓨팅 리소스가 가장 적습니다.

성능 비교#

레이어 동결의 효과를 확인하기 위해 공식 COCO 사전 학습 가중치(yolov5m.pt)에서 시작하여 Pascal VOC 데이터셋으로 YOLOv5m을 50 에포크 동안 학습했습니다. 세 가지 시나리오를 비교했습니다. 모든 레이어 학습(--freeze 0), 백본 동결(--freeze 10), 최종 탐지 레이어를 제외한 모든 레이어 동결(--freeze 24)입니다.

# Example command for training with backbone frozen
python train.py --batch 48 --weights yolov5m.pt --data voc.yaml --epochs 50 --cache --img 512 --hyp data/hyps/hyp.VOC.yaml --freeze 10

정확도 결과#

결과에 따르면 레이어를 동결하면 학습 속도가 크게 빨라질 수 있지만 최종 평균 정밀도 (mAP)가 약간 낮아질 수 있습니다. 일반적으로 모든 레이어를 학습하면 정확도가 가장 높고, 더 많은 레이어를 동결하면 잠재적으로 성능이 낮아지는 대신 학습 속도가 빨라집니다.

다양한 동결 전략을 비교한 학습 mAP50 결과 학습 중 mAP50 비교

다양한 동결 전략을 비교한 학습 mAP50-95 결과 학습 중 mAP50-95 비교

YOLOv5 frozen layer training performance

*Summary table of performance metrics*

리소스 사용량#

더 많은 레이어를 동결하면 GPU 메모리 요구 사항과 전체 사용률이 크게 감소합니다. 따라서 제한된 하드웨어 리소스에서 작업할 때 동결된 레이어를 사용한 전이 학습이 매력적인 선택이 되며, 그렇지 않으면 어려울 수 있는 더 큰 모델 또는 더 큰 이미지 크기를 사용하여 학습할 수 있습니다.

학습 중 할당된 GPU 메모리 비율 GPU 메모리 할당량 (%)

학습 중 GPU 메모리 사용률 GPU 사용률 (%)

레이어 동결을 사용해야 하는 경우#

전이 학습 중 레이어 동결은 다음과 같은 여러 상황에서 특히 유리합니다.

제한된 컴퓨팅 리소스: GPU 메모리 또는 처리 성능에 제약이 있는 경우입니다.
소규모 데이터셋: 대상 데이터셋이 원래 사전 학습 데이터셋보다 훨씬 작은 경우 레이어를 동결하면 과적합을 방지하는 데 도움이 됩니다.
신속한 프로토타이핑: 초기 평가를 위해 기존 모델을 새로운 작업이나 도메인에 빠르게 적용해야 하는 경우입니다.
유사한 특징 도메인: 새 데이터셋의 저수준 특징이 모델이 사전 학습된 데이터셋의 특징과 매우 유사한 경우입니다.

용어집 항목에서 전이 학습의 다양한 측면을 자세히 살펴보고, 성능을 최적화하기 위해 하이퍼파라미터 튜닝과 같은 기법을 고려해 보세요.

지원 환경#

Ultralytics는 CUDA, CuDNN, Python, PyTorch와 같은 필수 종속 항목이 사전 설치된 다양한 즉시 사용 가능한 환경을 제공합니다.

무료 GPU 노트북:
Google Cloud: GCP 빠른 시작 가이드
Amazon: AWS 빠른 시작 가이드
Azure: AzureML 빠른 시작 가이드
Docker: Docker 빠른 시작 가이드

프로젝트 상태#

이 배지는 모든 YOLOv5 GitHub Actions 지속적 통합 (CI) 테스트가 성공적으로 통과했음을 확인합니다. 이러한 CI 테스트는 학습, 검증, 추론, 내보내기, 벤치마크 등 주요 작업 전반에서 YOLOv5의 기능과 성능을 엄격하게 평가합니다. 또한 macOS, Windows, Ubuntu에서 일관되고 안정적으로 작동하는지 확인하며, 24시간마다 한 번씩 그리고 새로운 코드 커밋이 발생할 때마다 자동으로 실행됩니다.

기여자

RAraimbekovm² GLglenn-jocher² PDpderrenger¹

생성일 2개월 전업데이트됨 지난주