[python] 파이썬 PyTorch에서 특정 파라미터를 고정하여 모델을 학습하는 방법은?

22 Nov 2023

python

PyTorch에서는 모델을 학습할 때 특정 파라미터를 고정하는 기능을 제공합니다. 이 기능을 사용하면 특정 파라미터가 학습과정에서 업데이트되지 않고 고정된 상태로 유지됩니다. 이는 어떤 파라미터를 고정해야할 때 유용한 기능입니다. 파이썬 PyTorch를 사용하여 모델을 학습하는 방법을 알아보겠습니다.

먼저, 모델을 정의하고 학습 데이터를 준비합니다. 그런 다음, optimizer를 생성하고 학습을 수행하는 과정에서 특정 파라미터를 고정하려면 다음 단계를 따릅니다.

import torch
import torch.nn as nn
import torch.optim as optim

# 모델 정의
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.linear = nn.Linear(10, 1)
    
    def forward(self, x):
        return self.linear(x)

model = MyModel()

# 학습 데이터 준비
train_data = torch.randn(100, 10)
target = torch.randn(100, 1)

# optimizer 생성
optimizer = optim.SGD(model.parameters(), lr=0.01)

# 일부 파라미터 고정
parameters_to_freeze = [model.linear.weight]  # 예를 들어, linear layer의 weight 파라미터를 고정

for param in parameters_to_freeze:
    param.requires_grad = False

# 학습 수행
for epoch in range(10):
    optimizer.zero_grad()
    output = model(train_data)
    loss = nn.MSELoss()(output, target)
    loss.backward()
    optimizer.step()

위의 코드에서는 MyModel 클래스를 정의하고, nn.Linear를 이용하여 모델의 레이어를 생성했습니다. 그리고 학습 데이터인 train_data와 target을 준비했습니다.

model.parameters()를 통해 모델의 모든 파라미터를 얻을 수 있습니다. optimizer로 SGD를 사용하고, lr 파라미터로 learning rate를 설정했습니다.

parameters_to_freeze 변수에 고정하고자 하는 파라미터를 리스트로 추가합니다. 위의 코드에서는 model.linear.weight를 고정하고 있습니다. requires_grad를 False로 설정하여 해당 파라미터가 학습 과정에서 업데이트되지 않도록 만듭니다.

마지막으로, optimizer.step()을 호출하여 모델 파라미터를 업데이트합니다. 드롭아웃이나 배치 정규화와 같은 다른 파라미터들은 여전히 학습됩니다.

참고 문서로는 PyTorch 공식 문서를 참고하시면 더 자세한 내용을 확인하실 수 있습니다.