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

22 Nov 2023

python

아래는 예제 코드입니다:

import torch
import torch.nn as nn

# 모델 정의
class MyModel(nn.Module):
    def __init__(self):
        super(MyModel, self).__init__()
        self.layer1 = nn.Linear(10, 20)
        self.layer2 = nn.Linear(20, 5)

    def forward(self, x):
        x = self.layer1(x)
        x = self.layer2(x)
        return x

model = MyModel()

# 가중치 고정
model.layer1.weight.requires_grad = False

# 가중치가 고정되었는지 확인
for name, param in model.named_parameters():
    print(name, param.requires_grad)

# 학습을 위한 데이터와 loss 함수 설정 등 생략

# 학습 과정
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
for epoch in range(10):
    optimizer.zero_grad()
    output = model(input_data)
    loss = loss_fn(output, target)
    loss.backward()
    optimizer.step()

# 학습을 거치지 않을 계층의 가중치는 변하지 않습니다.

위 코드에서 model.layer1.weight의 requires_grad 속성을 False로 설정하여 해당 가중치가 고정되었습니다. 학습 과정에서 optimizer.step() 메서드가 호출될 때 가중치가 업데이트되지 않는 것을 확인할 수 있습니다.

자세한 내용은 공식 PyTorch 문서를 참조하시기 바랍니다.