PyTorch 튜토리얼 — 파인튜닝과 고정 특징 추출기

PyTorch 튜토리얼 — 파인튜닝과 고정 특징 추출기

PyTorch는 컴퓨터 비전 전이 학습 예제로 두 가지 전략을 보여줘요. ImageNet으로 학습한 ConvNet을 우리 데이터셋(개미·벌 분류)에 적용하는 흐름이에요.

전략 1: Finetuning the ConvNet

무작위 초기화 대신 사전 학습 가중치로 모델을 초기화하고, 그 상태 그대로 우리 데이터로 처음부터 학습한다. 전체 네트워크가 우리 작업을 반영해 조금씩 적응한다.

전략 2: ConvNet as fixed feature extractor

마지막 완전연결(FC) 레이어만 새로 바꾸고, 나머지 사전 학습된 파라미터는 고정(requires_grad=False)한다. 고정 부분은 특징 추출기로만 동작하고, 새 레이어만 학습해 가볍고 빠르다.

학습 루프 예시

model_conv = torchvision.models.resnet18(weights=IMAGENET1K_V1)
for param in model_conv.parameters():
    param.requires_grad = False
model_conv.fc = nn.Linear(512, 2)   # 새 분류 헤드

optim.lr_scheduler.StepLR로 7 에폭마다 학습률을 0.1배로 줄이며 훈련하고, 검증 손실이 가장 낮은 최고 모델을 체크포인트로 저장한다.

더 알아보기

• PyTorch 전이 학습 튜토리얼
• TensorFlow Keras 전이 학습 가이드