Imagenet으로 Pretraining된 모델을 Fine Tuning 해보는 것을보며 Fine Tuning에 대해 이해하기

ImageNet으로 사전 학습된 모델을 fine-tuning하는 방법은, 사전 학습된 모델을 새로운 데이터셋에 맞춰 조정하는 과정입니다. 사전 학습된 모델이 ImageNet과 같은 대형 데이터셋에서 이미 학습된 가중치를 사용하기 때문에, fine-tuning은 더 작은 데이터셋에서도 효율적으로 학습할 수 있도록 도와줍니다. 보통 다음과 같은 단계를 따릅니다:

1. 사전 학습된 모델 불러오기

먼저 ImageNet에서 사전 학습된 모델을 불러옵니다. 일반적으로 ResNet, VGG, Inception, EfficientNet 같은 모델들이 ImageNet으로 사전 학습되어 널리 사용됩니다. 프레임워크(예: PyTorch, TensorFlow)를 사용하면 사전 학습된 모델을 쉽게 불러올 수 있습니다.

from torchvision import models

# ImageNet에서 사전 학습된 ResNet50 불러오기
model = models.resnet50(pretrained=True)

2. 마지막 레이어 변경

ImageNet의 분류 작업은 1000개의 클래스에 대해 학습되었으므로, fine-tuning할 데이터셋에 맞게 마지막 레이어를 변경해야 합니다. 새로운 데이터셋이 더 적은 클래스나 다른 클래스 집합을 가지고 있다면, 기존의 마지막 레이어를 제거하고 새로운 레이어를 추가합니다.

예를 들어, ResNet50 모델의 마지막 레이어를 새로운 데이터셋에 맞춰 10개의 클래스로 바꿉니다.

import torch.nn as nn

# ResNet50의 마지막 레이어를 10개 클래스로 변경
model.fc = nn.Linear(model.fc.in_features, 10)

3. 동결(Freeze) 또는 부분 동결

모델의 초기 레이어들은 이미지의 기본적인 패턴(엣지, 색상 등)을 학습하는 반면, 뒤쪽 레이어들은 더 고수준의 특성을 학습합니다. 일반적으로 초기 레이어는 이미 잘 학습되었기 때문에, 이를 동결(freeze)하여 학습 중 업데이트되지 않도록 하고, 마지막 몇 개의 레이어만 학습되도록 설정합니다.

# 모든 레이어를 동결
for param in model.parameters():
    param.requires_grad = False

# 마지막 레이어만 학습 가능하도록 설정
for param in model.fc.parameters():
    param.requires_grad = True

또는, 특정 레이어까지만 동결하고 그 이후의 레이어들은 학습 가능하도록 할 수 있습니다. 이는 데이터셋의 크기와 복잡도에 따라 다릅니다.

4. 학습 설정

학습할 레이어를 지정했으면, 손실 함수와 옵티마이저를 설정하고 학습을 시작합니다.

• 손실 함수는 주로 CrossEntropyLoss를 사용합니다.
• 옵티마이저는 사전 학습된 가중치에 비해 새로운 레이어의 가중치를 더 빠르게 학습하도록 작은 학습률을 설정하는 것이 일반적입니다.

import torch.optim as optim

# 교차 엔트로피 손실 함수
criterion = nn.CrossEntropyLoss()

# 옵티마이저 - 모델의 학습 가능한 파라미터만 업데이트 (예: 마지막 레이어)
optimizer = optim.SGD(model.fc.parameters(), lr=0.001, momentum=0.9)

5. 학습

모델을 새로운 데이터셋으로 학습시킵니다. 보통 일반적인 학습 프로세스와 같으며, 데이터 증강(augmentation)을 추가하여 과적합을 방지할 수 있습니다.

for epoch in range(epochs):
    model.train()
    running_loss = 0.0
    for inputs, labels in dataloader:
        # 데이터 전처리 및 GPU 사용 가능시 이동
        inputs, labels = inputs.to(device), labels.to(device)

        # 옵티마이저 초기화
        optimizer.zero_grad()

        # 순전파
        outputs = model(inputs)
        loss = criterion(outputs, labels)

        # 역전파 및 옵티마이저 업데이트
        loss.backward()
        optimizer.step()

        running_loss += loss.item()

6. 전체 모델 미세 조정 (Optional)

초기에는 동결된 레이어만 학습했지만, 전체 모델을 fine-tune하고 싶다면 동결을 해제하고 모델의 모든 레이어에 대해 미세 조정을 할 수 있습니다. 이 경우 더 작은 학습률로 전체 모델을 학습시키는 것이 좋습니다.

# 모든 레이어의 동결을 해제하고 학습 가능하도록 설정
for param in model.parameters():
    param.requires_grad = True

# 전체 모델 학습을 위한 학습률을 더 작게 설정
optimizer = optim.SGD(model.parameters(), lr=0.0001, momentum=0.9)

7. 평가

학습 후에는 검증 세트나 테스트 세트에서 모델을 평가하여 성능을 확인합니다.

요약

ImageNet으로 사전 학습된 모델을 fine-tuning하는 일반적인 과정은 다음과 같습니다:
1. 사전 학습된 모델 불러오기
2. 마지막 레이어를 새로운 데이터셋에 맞게 변경
3. 필요한 레이어를 동결
4. 손실 함수 및 옵티마이저 설정
5. 새로운 데이터셋으로 학습
6. 필요한 경우 전체 모델에 대한 미세 조정

이를 통해 새로운 데이터셋에 맞춘 신속하고 효율적인 학습이 가능합니다.