[PyTorch] 모델 내 일부 레이어 freeze 하기

pretrained 모델을 가져와서 나의 데이터에 맞게 학습시킬 때 기존 모델의 전체적인 weight은 그대로 가져오고 classifier 역할을 하는 마지막 레이어만 내 데이터에 맞게 학습시키고 싶어졌다.

 

모델 내 레이어를 부분적으로 freeze 하기 위한 방법을 찾아보니 다음과 같이 모델의 파라미터를 불러와 requires_grad를 False로 두어 freeze하고 싶은 부분은 freeze하고, 가중치 업데이트를 하고 싶은 부분은 그대로 두면 된다!

나같은 경우는 모델의 마지막 레이어 'fc'만 학습을 시키고 싶어 아래와 같이 코드를 구성했다.

 

 

for name, param in model.named_parameters():
    if name.split('.')[1] == 'fc' :
        pass
    else :
        param.requires_grad = False

 

 

모델의 파라미터가 제대로 freeze 되었는지 확인하고 싶으면 아래와 같이 확인해보면 된다!

 

for name, param in model.named_parameters():
    print(name, param.requires_grad)