[python] 파이썬 PyTorch에서 모델의 가중치를 고정하여 특정 부분만 학습시키는 방법은?

22 Nov 2023

python

예를 들어, ResNet18 모델에서 마지막 fully connected layer를 고정하고 이전의 레이어만 학습시키고 싶다고 가정해보겠습니다.

import torch
import torchvision.models as models

# 미리 학습된 ResNet18 모델 가져오기
model = models.resnet18(pretrained=True)

# 모델의 모든 파라미터의 requires_grad를 False로 설정
for param in model.parameters():
    param.requires_grad = False

# 마지막 fully connected layer의 requires_grad를 True로 설정
model.fc.requires_grad = True

# 모델을 학습과 평가를 위해 GPU 또는 CPU로 이동
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = model.to(device)

# 학습을 위한 optimizer와 loss function 설정 등의 코드 작성

# 학습 및 평가 loop 등의 코드 작성

위의 코드에서는 먼저 ResNet18 모델을 불러와서 pretrained=True로 설정하여 미리 학습된 가중치를 가져옵니다. 그 다음, 모델의 모든 파라미터의 requires_grad를 False로 설정하여 모든 가중치를 고정합니다. 그리고 마지막 fully connected layer의 requires_grad를 True로 설정하여 해당 부분만 학습할 수 있도록 합니다. 마지막으로 모델을 GPU 또는 CPU로 이동시키고, 학습 및 평가를 위한 optimizer와 loss function을 설정한 후, 학습 및 평가 loop를 작성합니다.

이 방법을 사용하면 모델의 특정 부분만 학습하고 싶을 때 해당 부분의 가중치를 고정할 수 있습니다.