[pytorch] 모델의 일부 레이어 웨이트 업데이트 막기 | model freezing (모델 프리징)

PyTorch에서 모델 일부분의 학습을막기 위해서는 해당 부분의 매개변수에 대해 requires_grad 속성을 False로 설정하면 된다. 이를 통해 그 매개변수에 대한 gradient가 계산되지 않아 웨이트가 업데이트되지 않는다.

기본적으로 model.named_parameters() 메서드를 사용하여 name과 param을 추출하고 원하는 name의 레이어만 requires_grad를 변경할 수 있다.

 

모델 프리징 예시

- Resnet 예시

import torch
import torch.nn as nn
import torchvision.models as models

# ResNet-50 모델 불러오기 (pre-trained)
resnet50 = models.resnet50(pretrained=True)

# layer1, layer2, layer3의 매개변수에 대해 requires_grad를 False로 설정
for name, param in resnet50.named_parameters():
    if 'layer1' in name or 'layer2' in name or 'layer3' in name:
        param.requires_grad = False

• resnet50의 'layer1', 'layer2', 'layer3' 부분만 required_grad 를 False로 변경하는 예시

 

 

 

- Densenet 예시

model = models.densenet121(pretrained=True)
print(model.state_dict().keys()) # model 의 key값을 확인

for param in model.parameters(): # model의 모든 parameter 를 freeze
    param.requires_grad = False

for param in model.features.denseblock4.denselayer16.parameters(): # 내가 training 시키기 원하는 부분만 unfreeze
    param.requires_grad = True

model.classifier = nn.Linear(1024,num_of_class) # classifier 부분 재설정

• 모델 웨이트 전체 required_grad 를 False로 변경하고 학습하고 싶은 레이어만 True로 변경해서 사용