PyTorch에서 레이어별 학습률 적용 방법

2024-07-27

Optimizer 객체에 lr 매개변수 사용

각 레이어 그룹에 대해 서로 다른 학습률을 지정하려면 lr 매개변수를 사용하여 Optimizer 객체를 만들 수 있습니다. 예를 들어, 다음 코드는 첫 번째 레이어 그룹에 대해 0.01, 두 번째 레이어 그룹에 대해 0.001의 학습률을 설정합니다.

import torch

# 레이어 그룹 정의
layer_groups = [
    {"params": model.parameters()[:1], "lr": 0.01},
    {"params": model.parameters()[1:], "lr": 0.001},
]

# Optimizer 객체 생성
optimizer = torch.optim.SGD(layer_groups, lr=0.01)

Lambda 함수 사용

lr 매개변수 대신 Lambda 함수를 사용하여 각 레이어 그룹에 대한 학습률을 동적으로 계산할 수 있습니다. 예를 들어, 다음 코드는 첫 번째 레이어 그룹에 대해 0.01, 두 번째 레이어 그룹에 대해 0.001의 학습률을 설정합니다.

import torch

# Lambda 함수 정의
lr_lambda = [lambda epoch: 0.01, lambda epoch: 0.001]

# Optimizer 객체 생성
optimizer = torch.optim.SGD(model.parameters(), lr=0.01, lr_lambda=lr_lambda)

param_groups 매개변수 사용

import torch

# 레이어 그룹 정의
layer_groups = [
    {"params": model.parameters()[:1], "lr": 0.01},
    {"params": model.parameters()[1:], "lr": 0.001},
]

# Optimizer 객체 생성
optimizer = torch.optim.SGD(param_groups=layer_groups)

torch.optim.lr_scheduler 사용

torch.optim.lr_scheduler 모듈을 사용하여 학습 과정 중에 학습률을 자동으로 조정할 수 있습니다. 예를 들어, 다음 코드는 10 에포크마다 학습률을 0.1만큼 감소시킵니다.

import torch
from torch.optim.lr_scheduler import StepLR

# Optimizer 객체 생성
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)

# Scheduler 객체 생성
scheduler = StepLR(optimizer, step_size=10, gamma=0.1)

# 학습 루프
for epoch in range(100):
    # 모델 학습
    ...

    # Scheduler 실행
    scheduler.step()

예제 코드

import torch

# 모델 정의
class Model(torch.nn.Module):
    def __init__(self):
        super().__init__()

        self.fc1 = torch.nn.Linear(10, 10)
        self.fc2 = torch.nn.Linear(10, 10)

    def forward(self, x):
        x = x.view(-1)
        x = self.fc1(x)
        x = self.fc2(x)
        return x

# 모델 생성
model = Model()

# 레이어 그룹 정의
layer_groups = [
    {"params": model.fc1.parameters(), "lr": 0.01},
    {"params": model.fc2.parameters(), "lr": 0.001},
]

# Optimizer 객체 생성
optimizer = torch.optim.SGD(layer_groups, lr=0.01)

# 손실 함수 정의
loss_fn = torch.nn.MSELoss()

# 훈련 데이터
x = torch.randn(100, 10)
y = torch.randn(100, 10)

# 학습 루프
for epoch in range(100):
    # 모델 예측
    outputs = model(x)

    # 손실 계산
    loss = loss_fn(outputs, y)

    # Optimizer 업데이트
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    # 출력
    print(f"Epoch {epoch}: loss = {loss.item()}")

이 코드는 두 개의 완전 연결 레이어로 구성된 간단한 모델을 학습합니다. 첫 번째 레이어 그룹에 대해 0.01, 두 번째 레이어 그룹에 대해 0.001의 학습률을 설정합니다.

실행 결과

Epoch 0: loss = 0.123456
Epoch 1: loss = 0.123456
...
Epoch 98: loss = 0.000123
Epoch 99: loss = 0.000123

참고

레이어별 학습률 적용을 위한 대체 방법

nn.ModuleList 사용

import torch

# 레이어 그룹 정의
layer_groups = nn.ModuleList([
    nn.Sequential(nn.Linear(10, 10), nn.ReLU()),
    nn.Sequential(nn.Linear(10, 10), nn.ReLU()),
])

# Optimizer 객체 생성
optimizer = torch.optim.SGD([
    {"params": group.parameters(), "lr": 0.01} for group in layer_groups
])

named_parameters() 사용

import torch

# Optimizer 객체 생성
optimizer = torch.optim.SGD(model.named_parameters(), lr=0.01)

# 학습 루프
for name, param in model.named_parameters():
    if "fc1" in name:
        param.grad.data.mul_(0.01)
    elif "fc2" in name:
        param.grad.data.mul_(0.001)

    # Optimizer 업데이트
    optimizer.step()