[PyTorch] CNN 설계 11. Transfer Learning

1. Package load 
2. 데이터셋 다운로드 및 훈련, 검증, 테스트 데이터셋 구성 
3. 하이퍼파라미터 세팅 
4. Dataset 및 DataLoader 할당 
5. 네트워크 설계 
6. train, validation, test 함수 정의 
7. 모델 저장 함수 정의 
8. 모델 생성 및 Loss function, Optimizer 정의 
9. Training 
10. 저장된 모델 불러오기 및 test 
11. Transfer Learning

 

 

new_model = torchvision.models.resnet50(pretrained=True)

 

다만 우리가 앞서 직접 정의한 모델(SimpleCNN)에서는 입력 이미지의 크기를 120x120로 한 것에 반해, 방금 불러온 ResNet은 입력 이미지 크기를 최소한 224x224로 가정하고 학습된 모델입니다. 따라서 입력 데이터 크기만 수정하여 DataLoader를 다시 정의하도록 하겠습니다.

 

 

 

우리가 불러온 ImageNet에 학습된 ResNet은 1000개의 class를 구분하는 네트워크입니다. 즉, 마지막 FC layer의 출력 뉴런 수가 1000개인 것입니다. 우리는 2개의 class를 구분하는 네트워크를 원하기 때문에, 마지막 FC layer만 출력을 2로 바꿔주고 이 부분에 대해서만 학습을 추가로 진행하면 됩니다.

상황에 따라서는 우리가 이번에 하는 것처럼 마지막 FC layer만 학습을 진행하는 것이 아니라 전체 네트워크에 대해서 학습을 이어서 진행하는 경우도 있습니다. 이를 파라미터 fine tuning(미세 조정)이라고 합니다. ImageNet에는 다양한 동물 class도 포함이 되어있습니다. 즉, 우리가 불러온 ResNet은 강아지와 고양이 같은 동물에 대한 특징을 이미 어느정도 잘 추출하는 네트워크인 것입니다. 따라서 fine tuning이 굳이 필요하지 않습니다. 게다가, 이번 실습과 같이 적은 양의 데이터를 통해 모델을 학습시키는 상황에서 fine tuning을 진행하는 것은 overfitting의 가능성이 커지는 것이기 때문에 오히려 성능을 낮추는 결과를 가져올 수 있습니다.

 

이러한 내용을 아래 코드에 구현해두었습니다.

• nn.Module.parameters를 통해 모델 파라미터에 대한 iterator를 가져올 수 있습니다. 먼저, 이 iterator를 통해 for문 을 돌며 모든 파라미터에 대해 requires_grad를 False로 바꿔줍니다. 이렇게 하면 이 파라미터들에 대해서는 기울기가 계산되지 않기 때문에, 파라미터 업데이트가 되지 않습니다.
• 그 다음으로는 맨마지막 FC layer를 새로 정의해주는 것입니다. 그런데 문제는 마지막 FC layer에 어떻게 접근하느냐 입니다. 우리가 구현을 하지 않았기 때문에 마지막 layer의 입력 뉴런수가 어떻게 되는지도 모르는 상황입니다. 그러면 우리가 할 일은 구현된 코드를 보는 것입니다. torchvision에 우리가 불러온 ResNet이 구현되어 있습니다. 이곳에 가면 구현된 ResNet의 마지막 레이어의 변수명이 self.fc로 되어있음을 알 수 있습니다. 따라서 우리는 모델.fc와 같은 방식으로 이 layer에 접근할 수 있습니다. 이러한 방식으로 마지막 FC layer의 입력 뉴런수를 가져오고, 출력 뉴런 수는 우리의 문제에 맞게 2로 하여 마지막 layer를 수정할 수 있습니다.

 

for param in new_model.parameters():
    param.requires_grad = False

num_ftrs = new_model.fc.in_features
## 코드 시작 ##
new_model.fc = nn.Linear(num_ftrs, 2)    # 위의 설명 1. 을 참고하여 None을 채우세요.
criterion = nn.CrossEntropyLoss()       # 위의 설명 2. 를 참고하여 None을 채우세요.
new_model = new_model.to(device)
optimizer = nn.optim.Adam(new_model.parameters(), learning_rate)       # 위의 설명 3. 을 참고하여 None을 채우세요.
## 코드 종료 ##
val_every = 1
saved_dir = './saved/ResNet50'

 

 

model_path = './saved/ResNet50/best_model.pt'
# model_path = './saved/pretrained/ResNet50/best_model.pt' # 모델 학습을 끝까지 진행하지 않은 경우에 사용

## 코드 시작 ##
checkpoint = torch.load(model_path, map_location = device)    # 위의 설명 1. 을 참고하여 None을 채우세요.
state_dict = check_point['net']    # 위의 설명 2. 를 참고하여 None을 채우세요.
new_model.load_state_dict(state_dict)                 # 위의 설명 3. 을 참고하여 None을 채우세요.
## 코드 종료 ##

 

겨우 마지막 FC layer만 학습시켰음에도 불구하고 우리의 SimpleCNN보다 성능이 훨씬 좋은 것을 볼 수 있습니다.