[DL] Gradient Accumulation & 활용

Gradient Accumulation 소개

Gradient Accumulation은 메모리 제약 때문에 큰 배치 사이즈를 사용할 수 없을 때 사용되는 기법입니다. 이 방법은 미니 배치를 통해 얻은 그래디언트를 여러 스텝에 걸쳐 누적한 후, 일정한 누적이 이루어진 후에 모델의 가중치를 업데이트합니다.

Gradient Accumulation의 작동 원리

• 기본 프로세스: 일반적인 학습 과정에서는 배치 사이즈 만큼의 데이터를 통해 한 번의 forward pass와 backpropagation을 진행합니다.
• Accumulation Step: Gradient Accumulation에서는 각 미니 배치의 그래디언트를 바로 업데이트하지 않고, 지정된 accumulation step만큼 그래디언트를 누적합니다.
• 가중치 업데이트: 지정된 step만큼 그래디언트가 누적되면, 이를 이용해 한 번에 가중치를 업데이트합니다.

예제: Gradient Accumulation 구현

아래 예제에서는 PyTorch를 사용하여 Gradient Accumulation을 구현한 예시를 보여줍니다. 여기서는 배치 사이즈가 128이고, accumulation step이 4인 경우를 가정합니다.

import torch
from torch.optim import Adam

# 모델, 옵티마이저, 데이터 로더 정의
model = MyModel()
optimizer = Adam(model.parameters(), lr=0.001)
data_loader = MyDataLoader()

# Gradient Accumulation
accumulation_steps = 4
for i, (inputs, labels) in enumerate(data_loader):
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss = loss / accumulation_steps  # Normalize our loss (if averaged)
    loss.backward()

    # Perform optimization step after 'accumulation_steps' iterations
    if (i + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

Gradient Accumulation의 장단점

• 장점:
  • 큰 배치 사이즈를 시뮬레이션할 수 있어, 리소스가 제한적인 환경에서도 큰 배치의 이점을 활용할 수 있습니다.
  • 메모리 사용량을 줄이면서도 큰 배치 사이즈의 효과를 얻을 수 있습니다.
• 단점:
  • 전체 훈련 시간이 증가할 수 있습니다. 각 스텝마다 가중치를 업데이트하지 않고 누적하는 과정이 필요하기 때문입니다.

Gradient Accumulation은 특히 GPU 메모리가 제한적인 상황에서 유용합니다. 이 기법을 사용하면 메모리 제약으로 인한 학습의 질 저하를 방지하면서도, 큰 배치 사이즈의 장점을 활용할 수 있습니다.