학습률 웜업(Learning Rate WarmUp)에 대한 이해

소개

딥러닝 모델을 효과적으로 학습시키기 위해서는 적절한 학습률(learning rate)을 설정하는 것이 매우 중요합니다. 학습률은 모델이 가중치를 업데이트하는 속도를 조절하는 매개변수로, 너무 높거나 낮으면 학습이 제대로 이루어지지 않을 수 있습니다. 이 중에서도 학습률을 초기에 적절히 설정하는 것이 중요한데, 이를 위해 사용되는 기법 중 하나가 학습률 웜업(Learning Rate WarmUp)입니다.

---

학습률 웜업이란?

학습률 웜업은 딥러닝 모델의 초기 학습 단계에서 학습률을 서서히 증가시키는 기법입니다. 이는 모델이 초기에는 불안정하고 예측이 부정확할 수 있으며, 학습률을 너무 높게 설정하면 그러한 불안정성이 더욱 심해질 수 있기 때문에 사용됩니다. 학습률 웜업을 통해 모델이 초기에 안정성을 유지하면서도 학습 속도를 향상시킬 수 있습니다.

---

학습률 웜업의 구현

학습률 웜업은 일반적으로 학습의 초기 몇 epoch 동안 적용됩니다. 이 기간 동안 학습률을 서서히 증가시키고, 이후에는 고정된 학습률을 사용하여 학습을 계속 진행합니다. 구체적으로는 다음과 같은 단계로 구현될 수 있습니다.

1. 모델과 최적화 알고리즘 선택: 학습률 웜업을 적용할 모델과 최적화 알고리즘을 선택합니다. 대부분의 딥러닝 프레임워크에서 제공하는 최적화 알고리즘은 학습률을 조절하는 매개변수를 가지고 있습니다.
2. 초기 학습률 설정: 학습률 웜업을 위해 초기 학습률을 설정합니다. 일반적으로는 학습의 초기 단계에서 사용할 학습률보다 낮은 값을 선택합니다.
3. 학습률 스케줄링: 초기 몇 epoch 동안 학습률을 서서히 증가시키는 스케줄을 설정합니다. 예를 들어, 선형적으로 학습률을 증가시키거나 지수적으로 증가시킬 수 있습니다.
4. 학습률 업데이트: 각 epoch 또는 일정한 주기마다 학습률을 업데이트합니다. 이 때, 설정한 학습률 스케줄에 따라 증가된 학습률을 사용합니다.

---

학습률 웜업의 구현 - PyTorch + transformer module

아래 예시는 PyTorch와 Transformer 모듈의 get_cosine_schedule_with_warmup 함수를 이용하는 방법입니다. get_cosine_schedule_with_warmup의 원형은 아래와 같습니다.

위 예시와 같이 지정한 Learning Rate까지 특정 Epoch만큼 빠르게 도달한 후 다음 Epoch마다 감소하는 Cosine 그래프 형태를 띄면서 learning rate를 변화시킵니다.

import torch.nn as nn
from transformers import get_cosine_schedule_with_warmup

criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=0.00006, weight_decay = 0.0001)

epochs = 39

scheduler = get_cosine_schedule_with_warmup(
                optimizer,
                num_warmup_steps=len(loader_train) * 3,
                num_training_steps=len(loader_train) * epochs
            )

from sklearn.metrics import roc_auc_score
from tqdm.notebook import tqdm

for epoch in range(epochs):
    model.train()
    epoch_train_loss = 0

    for images, labels in tqdm(loader_train):
        # ... Train model
        # Step schduler for each epoch.
        scheduler.step()

    # ... Validate model
    with torch.no_grad():
        ...

---

학습률 웜업의 구현 - Keras

TensorFlow의 Keras를 사용하여 cosine 학습률 웜업을 구현하는 방법은 사용자 정의 콜백(Callback)을 통해 할 수 있습니다. 아래는 cosine 학습률 웜업을 구현한 콜백 클래스의 예시 코드입니다.

import tensorflow as tf
from tensorflow.keras.callbacks import Callback
from tensorflow.keras import backend as K

class CosineAnnealingWarmUpLR(Callback):
    def __init__(self, initial_lr, warmup_epochs, max_epochs, steps_per_epoch, min_lr=0):
        super(CosineAnnealingWarmUpLR, self).__init__()
        self.initial_lr = initial_lr
        self.warmup_epochs = warmup_epochs
        self.max_epochs = max_epochs
        self.steps_per_epoch = steps_per_epoch
        self.min_lr = min_lr
        self.history = {}

    def on_epoch_begin(self, epoch, logs=None):
        if epoch < self.warmup_epochs:
            lr = (self.initial_lr - self.min_lr) * (epoch / self.warmup_epochs) + self.min_lr
        else:
            lr = self.min_lr + 0.5 * (self.initial_lr - self.min_lr) * (
                        1 + np.cos(((epoch - self.warmup_epochs) / (self.max_epochs - self.warmup_epochs)) * np.pi))

        K.set_value(self.model.optimizer.lr, lr)
        self.history.setdefault('lr', []).append(K.get_value(self.model.optimizer.lr))

    def on_epoch_end(self, epoch, logs=None):
        logs = logs or {}
        logs['lr'] = K.get_value(self.model.optimizer.lr)
        self.history.setdefault('lr', []).append(logs['lr'])

위 코드에서 CosineAnnealingWarmUpLR 클래스는 초기 학습률, 웜업 에폭 수, 최대 에폭 수, 에폭 당 스텝 수 및 최소 학습률을 인자로 받습니다. 각 에폭의 시작 시에 학습률을 계산하고 설정하며, 이를 통해 cosine 학습률 스케줄링을 구현하고 웜업을 수행합니다.

이 콜백 클래스를 모델 훈련 시에 사용하여 cosine 학습률 웜업을 적용할 수 있습니다.

# 모델 컴파일
model.compile(optimizer='adam', loss='sparse_categorical_crossentropy', metrics=['accuracy'])

# 콜백 초기화
warmup_epochs = 5
max_epochs = 20
steps_per_epoch = len(train_data) // batch_size  # train_data는 데이터셋 객체입니다.

# 콜백 생성
cosine_warmup_lr = CosineAnnealingWarmUpLR(initial_lr=0.001, warmup_epochs=warmup_epochs, max_epochs=max_epochs, steps_per_epoch=steps_per_epoch)

# 모델 훈련
history = model.fit(train_data, epochs=max_epochs, callbacks=[cosine_warmup_lr])

위와 같이 모델 훈련 시에 CosineAnnealingWarmUpLR 콜백을 추가하여 cosine 학습률 웜업을 적용할 수 있습니다.

학습률 웜업의 장점

• 모델의 초기 학습 과정을 안정화시킬 수 있습니다.
• 학습률을 너무 높게 설정할 경우 발생할 수 있는 발산 문제를 방지할 수 있습니다.
• 초기에는 불안정한 예측으로 인한 손실을 최소화하고, 안정적인 학습을 돕습니다.

---

마무리

딥러닝에서의 학습률 웜업은 모델의 초기 학습 과정을 안정화시키고 더 나은 성능을 얻기 위한 중요한 기법 중 하나입니다. 적절한 학습률 웜업을 통해 모델이 초기에 안정적으로 학습되도록 보장할 수 있으며, 이는 전체 학습 과정의 성능 향상에 큰 영향을 미칠 수 있습니다.

---

Referenece

• https://huggingface.co/docs/transformers/main_classes/optimizer_schedules