[python] 파이썬 PyTorch에서 학습률 스케줄링을 수행하는 방법은?

22 Nov 2023

python

학습률 스케줄링은 딥 러닝 모델의 학습 과정에서 학습률(learning rate)을 조절하는 방법입니다. 적절한 학습률 스케줄링은 모델의 수렴을 높이고, 학습 속도를 개선할 수 있습니다. PyTorch에서는 다양한 학습률 스케줄러를 제공하고 있으며, 간단한 코드로 적용할 수 있습니다.

여기에는 PyTorch에서 자주 사용되는 두 가지 학습률 스케줄링 방법을 소개하겠습니다.

StepLR: StepLR은 주어진 step size마다 주어진 gamma로 학습률을 감소시키는 방법입니다. 즉, 주어진 step size마다 학습률을 gamma만큼 감소시킵니다.

import torch
import torch.optim as optim
import torch.nn as nn
from torch.optim import lr_scheduler

# 모델 정의
model = ...

# 손실 함수 및 옵티마이저 정의
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 스케줄러 정의
scheduler = lr_scheduler.StepLR(optimizer, step_size=30, gamma=0.1)

# 학습 루프
for epoch in range(num_epochs):
    # ...
    # 일반 학습 코드
    # ...
    
    # 스케줄러로 학습률 감소
    scheduler.step()

CosineAnnealingLR: CosineAnnealingLR은 cosine 함수 형태로 학습률을 감소시키는 방법입니다. 주어진 T_max에 도달할 때까지 학습률이 감소하고, 이후 다시 증가하는 패턴을 가집니다.

import torch
import torch.optim as optim
import torch.nn as nn
from torch.optim import lr_scheduler

# 모델 정의
model = ...

# 손실 함수 및 옵티마이저 정의
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.1)

# 스케줄러 정의
scheduler = lr_scheduler.CosineAnnealingLR(optimizer, T_max=100)

# 학습 루프
for epoch in range(num_epochs):
    # ...
    # 일반 학습 코드
    # ...
    
    # 스케줄러로 학습률 감소
    scheduler.step()

이렇게 간단한 코드로 PyTorch에서 학습률 스케줄링을 적용할 수 있습니다. 추가적으로 다른 스케줄링 방법이나 매개 변수에 대해 자세한 정보를 알고 싶다면 PyTorch 공식 문서를 참고하시기 바랍니다.

PyTorch 공식 문서: Learning Rate Scheduling