반려동물 품종 분류기 만들기 (5) Transfer learning: ResNet fine-tuning
사전학습 ResNet-34의 분류층만 학습하는 linear probing과 전체를 낮은 learning rate로 학습하는 full fine-tuning을 순서대로 실험하고, baseline 대비 성능 차이를 기록했습니다.
사전학습 모델 불러오기
backbone은 ResNet-34를 사용합니다. 딥러닝 기초 10편에서 구조를 다룬 모델이고, 이 데이터 규모에는 ResNet-50보다 작고 빠른 쪽이 적당합니다.
1
2
3
4
from torchvision.models import resnet34, ResNet34_Weights
model = resnet34(weights=ResNet34_Weights.IMAGENET1K_V1)
model.fc = nn.Linear(model.fc.in_features, 37)
ImageNet 1000 클래스용 마지막 fc 층을 37 클래스용 새 Linear로 교체합니다. 이 층만 무작위 초기화 상태이고, 나머지는 전부 ImageNet에서 배운 가중치입니다. 3편에서 normalize를 ImageNet 통계로 맞춰둔 것이 여기서 전제가 됩니다. 사전학습 당시와 같은 분포의 입력이 들어가야 배운 특징이 그대로 동작합니다.
단계 1: Linear probing
먼저 backbone을 전부 얼리고 새 fc 층만 학습합니다.
1
2
3
4
5
6
for param in model.parameters():
param.requires_grad = False
for param in model.fc.parameters():
param.requires_grad = True
optimizer = torch.optim.AdamW(model.fc.parameters(), lr=1e-3)
학습되는 파라미터가 fc 층의 약 2만 개뿐이라 epoch당 시간이 짧고, 5 epoch 만에 validation accuracy가 86.7%에 도달했습니다. baseline의 41.3%에서 45%p가 오른 것입니다. 모델 구조를 바꾼 것도, 데이터를 늘린 것도 아니고, ImageNet에서 배운 특징 추출기를 재사용했을 뿐입니다.
이 결과가 말해주는 것은 분명합니다. 귀 모양이나 털 결 같은 시각 특징은 이미 backbone 안에 있고, 우리 문제에는 그 특징을 37개 품종으로 묶는 분류층만 새로 필요했던 것입니다.
단계 2: Full fine-tuning
다음으로 backbone까지 전부 풀고 학습합니다. 여기서 중요한 것은 learning rate입니다. 새 fc 층에는 1e-3을 그대로 쓰되, 이미 잘 학습된 backbone에는 10분의 1인 1e-4를 적용합니다.
1
2
3
4
5
6
7
8
9
for param in model.parameters():
param.requires_grad = True
optimizer = torch.optim.AdamW([
{"params": model.fc.parameters(), "lr": 1e-3},
{"params": [p for n, p in model.named_parameters()
if not n.startswith("fc")], "lr": 1e-4},
])
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=10)
backbone에 큰 learning rate를 적용하면 무작위 초기화된 fc 층에서 내려오는 큰 gradient가 사전학습 가중치를 망가뜨립니다. 실제로 전체를 1e-3으로 풀었더니 첫 epoch에서 validation이 60%대로 떨어졌다가 회복하지 못했습니다. 낮은 learning rate로 10 epoch 학습한 결과는 validation accuracy 93.1%입니다.
세 모델 비교
동일한 분할과 전처리에서의 결과입니다. 학습 시간은 Colab T4 기준입니다.
| 모델 | validation accuracy | 학습 시간 |
|---|---|---|
| SmallCNN (baseline) | 41.3% | 약 25분 (30 epoch) |
| ResNet-34 linear probing | 86.7% | 약 8분 (5 epoch) |
| ResNet-34 full fine-tuning | 93.1% | 약 30분 (10 epoch) |
- baseline에서 linear probing으로 45.4%p 상승. 사전학습 특징 재사용의 효과입니다
- probing에서 full fine-tuning으로 6.4%p 추가 상승. ImageNet 특징을 품종 구분에 맞게 미세 조정한 효과입니다
- probing은 가장 짧은 시간에 86.7%를 냅니다. 빠르게 기준점이 필요할 때 probing부터 하는 것이 좋은 순서라는 것도 확인했습니다
딥러닝 기초 11편의 “작은 데이터셋일수록 fine-tuning이 유리하다”는 명제가 이 문제에서는 51.8%p 차이로 측정되었습니다.
남은 질문
validation 93.1%가 최종 성적은 아닙니다. test 셋은 아직 열지 않았고, 6.9%의 오답이 어떤 품종에서 나오는지도 아직 모릅니다. 다음 글에서 test 셋으로 최종 평가를 하고, confusion matrix와 Grad-CAM으로 모델이 어디서 왜 틀리는지 분석합니다.