PyTorch의 전이 학습(Transfer Learning)에서 Freezing 여부에 따른 성능 차이 및 유의점

  파이토치(PyTorch)의 공식 문서에서 전이 학습(transfer learning)에 관해 설명하고 있는 문서는 다음과 같다.

 

  ▶ PyTorch Transfer Learning: pytorch.org/tutorials/beginner/transfer_learning_tutorial.html

Transfer Learning for Computer Vision Tutorial — PyTorch Tutorials 1.7.1 documentation

  공식 문서에서는 전이 학습(transfer learning)의 대표적인 두 가지 시나리오를 언급한다.

 

  1. 전체 네트워크를 fine-tuning 하는 방식

  2. 사전학습된(pre-trained) 네트워크를 고정된 특징 추출기(fixed feature extractor)로 사용하는 방식

 

  여기에서 1번과 2번의 실제 구현상의 차이점은 사전학습된 네트워크에 대하여 다음의 코드 부분을 넣느냐 마느냐이다. 아래 코드는 사전학습된(pre-trained) 네트워크의 가중치를 고정할 때 사용하는 코드이다.

 

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

 

  만약 2번 방식(fixed feature extractor)대로 사전학습된 네트워크의 가중치를 특징 추출기로 고정한다면, 뒤쪽에 있는 FC 레이어만 업데이트가 될 것이다. 또한 이 경우에는 앞쪽 레이어에 대한 기울기(gradient)를 계산하지 않아도 되기 때문에 학습 속도가 빨라진다. 참고로 2번의 방식을 사용하는 경우 optimizer에서는 FC 레이어의 파라미터에 대해서만 업데이트한다고 명시해야 한다. (optimizer는 계산된 기울기(gradient)를 이용해 업데이트(update)를 수행한다.)

 

  [참고 1] requires_grad를 False로 설정한 레이어에 대하여 optimizer에서 업데이트를 하겠다고 명시하더라도, 어차피 구해진 gradient 값 자체가 없기 때문에 업데이트가 수행되지 않기는 한다.

 

  [참고 2] PyTorch에서 모델을 초기화하면 기본 설정으로 requires_grad 값이 True가 된다. (이는 사전 학습된 네트워크를 불러올 때에도 마찬가지다!) 따라서 별도로 명시하지 않는다면 자동으로 기울기(gradient)를 추적하기 때문에, 학습하지 않고자 하는 레이어에 대해서는 requires_grad 값을 명시적으로 False로 설정할 필요가 있다.

 

  다만 공식 문서에서는 간단한 이진 분류(binary classification) 예시를 들고 있기 때문에 1번과 2번의 성능 차이가 크게 나지 않으며, 디테일한 하이퍼 파라미터 세팅을 하지 않아도 높은 성능이 나온다. 하지만 실제로 CIFAR-10과 같이 클래스의 개수가 많은 데이터셋을 이용하는 경우에는 성능 차이가 크게 날 수 있다.

 

  필자의 경우 클래스가 3개인 경우, 클래스가 10개인 경우에 대하여 학습을 진행해 보았다. 이때 2번 방법대로 마지막 FC 레이어만 학습하도록 한 경우에는 학습이 정상적으로 수행되지 않았다. 하이퍼 파라미터 세팅에 많은 신경을 써야 하는 것으로 보인다. 따라서 클래스의 개수가 많은 상황에서 빠르게 높은 정확도(high accuracy)를 얻고 싶다면 1번의 방법대로 앞쪽 네트워크를 고정하지 않고 전체 네트워크를 fine-tuning 하는 것이 유리할 수 있다.

 

  예를 들어 필자의 경우 CIFAR-10에 대하여 전이 학습(transfer learning)을 수행한 경험이 있는데, 다른 코드 부분은 완전히 동일하게 유지한 상태로 한 번은 다음과 같은 코드를 사용했다.

 

net = torchvision.models.resnet18(pretrained=True)

# 마지막 레이어의 차원을 10차원으로 조절
num_features = net.fc.in_features
net.fc = nn.Linear(num_features, 10)
net = net.to(device)

 

  그리고 한 번은 다음과 같은 코드를 사용했다.

 

net = torchvision.models.resnet18(pretrained=True)
for param in net.parameters():
    param.requires_grad = False

# 마지막 레이어의 차원을 10차원으로 조절
num_features = net.fc.in_features
net.fc = nn.Linear(num_features, 10)
net = net.to(device)

 

  첫째 경우(fine-tuning)에는 한 번의 epoch만으로 순식간에 94% 정도의 test accuracy를 얻을 수 있었지만, 둘째 경우(fixed feature extractor)에는 여러 번의 epoch을 반복해도 90% 이상의 성능은 얻을 수 없었다. 다시 말해 2번 방법이 학습 속도 측면에서 유리할 수 있으나, 클래스가 많은 상황에서는 성능이 낮게 나오는 문제가 발생할 수 있다.

 

  [참고] 이 문제는 해외 기술 블로그에서도 자주 다루어지고 있는 내용이다. 정리하자면 다음과 같다.

 

  ① Frozen: 앞쪽의 특징 추출기(feature extractor)에 대하여 역전파를 수행하지 않는 방법이다. 일반적으로 목표 작업(target task)의 레이블(label) 수가 적고 오버피팅(overfitting)을 예방하기 위한 목적으로 사용된다.

  ② Fine-tuning: 앞쪽의 특징 추출기(feature extractor)에 대하여 역전파를 수행하는 방법이다. 일반적으로 목표 작업(target task)의 레이블(label) 수가 많을 때 사용한다.