[Deep Learning] Partial Fine-Tuning 해보기

Feature Extractor만 활용해보는 ResNet50 전이학습

 

딥러닝 모델을 학습할 때, 모델 전체를 처음부터 학습(From Scratch) 하는 것은 시간도 오래 걸리고, 많은 데이터가 필요하다.

그래서 우리는 보통 전이학습(Transfer Learning) 을 이용한다. 그중에서도 이번 글에서는 Partial Fine-Tuning, 즉 Feature Extractor만 사용하는 전이학습을 직접 실험해본다.

 

Fine-Tuning의 세 가지 방식

먼저 개념을 간단히 정리해보자.

From Scratch	랜덤 초기화로 처음부터 학습
Full Fine-Tuning	ImageNet 등 사전학습 가중치로 초기화 후 전체 재학습
Partial Fine-Tuning (Feature Extractor)	사전학습된 백본을 고정(Freeze)하고 마지막 FC 레이어만 학습

Partial Fine-Tuning은 모델이 이미 학습한 일반적인 시각적 특징(feature) 을 그대로 활용하고, 새 데이터셋에 맞게 마지막 분류기(classifier) 만 조정하는 방식이다.

 

코드 구현

import torch
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torchvision import models
from torch.utils.data import DataLoader
import torch.nn as nn
import torch.optim as optim
from torch.optim import lr_scheduler
import os, time, copy

# 학습 함수 정의
def train_resnet(model, criterion, optimizer, scheduler, num_epochs=25):
    best_model_wts = copy.deepcopy(model.state_dict())
    best_acc = 0.0
    for epoch in range(num_epochs):
        print(f'-------------- epoch {epoch+1} ----------------')

 

        for phase in ['train', 'val']:
            model.train() if phase == 'train' else model.eval()
            running_loss, running_corrects = 0.0, 0

            for inputs, labels in dataloaders[phase]:
                inputs, labels = inputs.to(DEVICE), labels.to(DEVICE)
                
                optimizer.zero_grad()

                with torch.set_grad_enabled(phase == 'train'):
                    outputs = model(inputs)
                    _, preds = torch.max(outputs, 1)
                    loss = criterion(outputs, labels)
                    if phase == 'train':
                        loss.backward()
                        optimizer.step()

                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            if phase == 'train':
                scheduler.step()

            epoch_loss = running_loss / dataset_sizes[phase]
            epoch_acc = running_corrects.double() / dataset_sizes[phase]
            print(f'{phase} Loss: {epoch_loss:.4f} Acc: {epoch_acc:.4f}')

            if phase == 'val' and epoch_acc > best_acc:
                best_acc = epoch_acc
                best_model_wts = copy.deepcopy(model.state_dict())

        print(f'Best val Acc so far: {best_acc:.4f}')
    model.load_state_dict(best_model_wts)
    return model

 

 

• model.train() => train 단계: 모델이 가중치를 업데이트함
• model.eval()  => val 단계: 학습된 모델의 일반화 성능을 확인함

 

optimizer.zero_grad() → 이전 배치에서 계산된 gradient 초기화

=> PyTorch는 기본적으로 gradient를 누적하므로 매 step마다 초기화

 

with torch.set_grad_enabled(phase == 'train'):
    outputs = model(inputs)
    _, preds = torch.max(outputs, 1)
    loss = criterion(outputs, labels)
    if phase == 'train':
        loss.backward()
        optimizer.step()

 

• torch.set_grad_enabled(phase == 'train')
  → 학습 단계(train)에서만 gradient 계산 활성화,
  검증(val)에서는 비활성화 → 메모리 절약 및 속도 향상
• outputs = model(inputs) → forward propagation 수행
• torch.max(outputs, 1) → softmax 결과 중 가장 높은 확률을 가진 클래스의 index를 예측(preds)
• criterion(outputs, labels) → CrossEntropyLoss 계산
• loss.backward() → gradient 계산 (오차 역전파)
• optimizer.step() → 가중치 갱신

if phase == 'train':
    scheduler.step()

 

• 학습 단계가 끝날 때마다 StepLR 스케줄러를 한 번 호출
• 일정 에폭(step_size)마다 학습률을 gamma 비율만큼 줄임

data_transforms = {
    'train': transforms.Compose([
        transforms.Resize([64, 64]),          # 이미지 크기 통일
        transforms.RandomHorizontalFlip(),    # 좌우 반전
        transforms.RandomVerticalFlip(),      # 상하 반전
        transforms.RandomCrop(52),            # 무작위 크롭 (데이터 증강)
        transforms.ToTensor(),                # Tensor로 변환
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])  # ImageNet 정규화 값
    ]),
    'val': transforms.Compose([
        transforms.Resize([64, 64]),
        transforms.RandomCrop(52),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406],
                             [0.229, 0.224, 0.225])
    ])
}

 

Normalize()는 ImageNet 데이터셋의 평균, 표준편차로 정규화하여 사전학습된 모델(ResNet50) 의 입력 분포와 맞춰준다.

 

데이터셋 및 데이터로더 정의

data_dir = './splitted'  # train, val 폴더가 들어있는 상위 폴더
image_datasets = {
    x: ImageFolder(root=os.path.join(data_dir, x),
                   transform=data_transforms[x])
    for x in ['train', 'val']
}

dataloaders = {
    x: DataLoader(image_datasets[x],
                  batch_size=BATCH_SIZE,
                  shuffle=True,
                  num_workers=4)   # CPU 4개로 데이터 병렬 로드
    for x in ['train', 'val']
}

dataset_sizes = {x: len(image_datasets[x]) for x in ['train', 'val']}
class_names = image_datasets['train'].classes

 

ImageFolder : 폴더 구조를 자동으로 라벨링하여 Dataset으로 변환

 

ResNet50 모델 정의 ( From Scratch, Full Fine-Tuning, Partial Fine-Tuning )

From Scratch

resnet = models.resnet50(pretrained=False)

num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 33)

resnet = resnet.to(DEVICE)  # GPU로 이동

 

완전 처음부터 학습

in_features : 기존 fc 레이어의 입력 뉴런 수(in_features) 를 가져오고

resnet.fc : 기존 1000-class 출력을 33-class로 교체한다

 

Full Fine-Tuning

resnet = models.resnet50(pretrained=True)
num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 33)

resnet = resnet.to(DEVICE)  # GPU로 이동

 

pretrained=True로 ImageNet 사전학습 가중치를 불러와서 train

 

Partial Fine-Tuning

resnet = models.resnet50(pretrained=True)
print(">>> Using Partial Fine-Tuning version (Feature Extractor)")

for param in resnet.parameters():
    param.requires_grad = False  # Feature Extractor 부분 Freeze

num_ftrs = resnet.fc.in_features
resnet.fc = nn.Linear(num_ftrs, 33)

for param in resnet.fc.parameters():
    param.requires_grad = True   # 마지막 FC만 학습

resnet = resnet.to(DEVICE)  # GPU로 이동

 

• pretrained=True로 ImageNet 사전학습 가중치를 불러옵니다.
• requires_grad=False → ResNet의 백본(Feature Extractor) 부분을 동결시켜 학습 제외
• requires_grad=True : 새로운 FC 레이어만 학습 대상

 

손실함수 / 옵티마이저 / 스케줄러 설정

criterion = nn.CrossEntropyLoss()  # 다중 클래스 분류용 Loss

optimizer_ft = optim.Adam(
    filter(lambda p: p.requires_grad, resnet.parameters()),
    lr=0.001
)

exp_lr_scheduler = lr_scheduler.StepLR(optimizer_ft, step_size=7, gamma=0.1)

 

결과

 

From Scratch는 가중치를 처음부터 학습했음에도 95.61%로 높은 성능을 보였다.

Full Fine-Tuning은 사전학습된 가중치를 기반으로 모든 레이어를 재학습해 99.13%의 최고 정확도를 기록했다.
반면 예상외로 Partial Fine-Tuning은 Feature Extractor를 고정하고 FC 레이어만 학습해 79.19%로 가장 낮은 정확도를 보였다.

데이터셋의 이미지 특성이 ImageNet과 달라서 그런 것 같기도 하고 모델이 새로운 데이터셋의 특성에 충분히 적응하지 못했던 것 같다.

 

검증 및 결과

import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision.transforms as transforms
from torchvision.datasets import ImageFolder
from torch.utils.data import DataLoader

def evaluate(model, test_loader):
    model.eval()
    test_loss = 0
    correct = 0

    with torch.no_grad():
        for data, target in test_loader:
            data, target = data.to(DEVICE), target.to(DEVICE)
            output = model(data)

            test_loss += F.cross_entropy(output, target, reduction='sum').item()

            pred = output.max(1, keepdim=True)[1]
            correct += pred.eq(target.view_as(pred)).sum().item()

    test_loss /= len(test_loader.dataset)
    test_accuracy = 100. * correct / len(test_loader.dataset)
    return test_loss, test_accuracy

if __name__ == '__main__':
    USE_CUDA = torch.cuda.is_available()
    DEVICE = torch.device("cuda" if USE_CUDA else "cpu")
    BATCH_SIZE = 256

    transform_resNet = transforms.Compose([
        transforms.Resize([64, 64]),
        transforms.ToTensor(),
        transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])

    # Evaluate model
    test_resNet = ImageFolder(root='./splitted/test', transform=transform_resNet)
    test_loader_resNet = torch.utils.data.DataLoader(test_resNet, batch_size=BATCH_SIZE, shuffle=False, num_workers=4)

    '''
    Compare the transfer-learned model with the learned model from scratch
    '''
    resnet = torch.load('resnet50_from_partial_0.7919.pt', weights_only=False)   # change model
    # resnet = torch.load('resnet50_from_pretrained_0.9913.pt') # change model
    # resnet = torch.load('resnet50_from_scratch_0.9561.pt') # change model
    print(resnet)
    resnet.to(DEVICE)
    test_loss, test_accuracy = evaluate(resnet, test_loader_resNet)

    print('test acc: ', test_accuracy)

 

ResNet50 Train과 Test 데이터셋으로 검증 결과 

From Scratch	랜덤 초기화	전체 레이어	약 95%	89.5%
Full Fine-Tuning	ImageNet pretrained	전체 레이어 재학습	약 99%	98.5%
Partial Fine-Tuning	ImageNet pretrained	FC layer만 재학습 (feature extractor 고정)	약 79%	73.3%

 

• Full Fine-Tuning은 사전학습된 가중치를 기반으로 전체를 재학습하여 가장 높은 정확도를 기록했다.
  → 사전학습 모델의 일반적 시각 특징을 유지하면서, 새로운 도메인에 맞게 최적화되었기 때문.
• From Scratch는 ImageNet 사전학습을 사용하지 않고도 약 89%의 테스트 정확도를 달성했지만,
  수렴 속도가 느리고 많은 데이터가 필요했다.
• Partial Fine-Tuning은 Feature Extractor를 고정했기 때문에
  새로운 도메인의 세부적인 질감이나 색상 패턴을 학습하지 못해 성능이 20% 이상 낮게 나타났다.
  즉, 사전학습된 일반적인 특성만으로는 도메인 적응이 어렵다는 한계가 드러났다.

참고자료

 

https://hi-ai0913.tistory.com/32

[딥러닝] 전이학습(Transfer learning)과 파인튜닝(Fine tuning)

 

2025.06.14 - [Technology Notes] - [Deep Learning] Transfer Learning과 Knowledge distillation

[Deep Learning] Transfer Learning과 Knowledge distillation