[PyTorch] pretrained model을 freeze하기, 선택적으로 freeze하기

* 모델 전체 레이어의 가중치를 freeze하기

model = torchvision.models.resnet18(pretrained=True) # 예시로 resnet18모델을 불러옴

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

* 선택적으로 freeze하기

ex. 첫 5개 레이어에 대해서만 가중치 freeze하기

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

count = 0

for child in model.children():
    if count < 5:
    	for param in child.parameters():
        	param.requires_grad = False
    count += 1

 

*참고 링크

pytorch에서 특정 layer freeze 하기 (학습하지 않기) freezing