Gradient accumulation (그레디언트 누적)

요약

---

Gradient Accumulation은 GPU 메모리 한계를 극복하면서 더 큰 배치(batch) 크기로 학습한 것과 동일한 효과를 내기 위해 사용하는 기법입니다.

 

문제: GPU 메모리 한계로 큰 배치를 한번에 학습할 수 없는 경우 발생

해결방법

1. 작은 배치에서 계산된 loss을 gradient을 누적합니다. (=.backward()만 호출하면 gradient값이 더해집니다.)
2. 정해진 횟수(gradient accumulation step)만큼 누적되었다면 optimizer로 가중키를 업데이트 합니다.
3. 최종기울기가 큰 배치에서 한번에 처리한 결과와 동일하도록 각 loss을 gradient accumulation step만큼 나눠주어 스케일링합니다.

 

코드 스니펫

---

아래와 같은 코드스니펫으로 사용될 수 있습니다. 주요 코드 설명은 아래와 같습니다.

• loss = loss / accumulation_steps : loss을 누적하여 이후에 backward을 할거기에 미리 accumulation_steps수로 나눠줍니다.
• loss.backward(): backward을 호출하면 gradient가 이전의 gradient값과 내부적으로 더해(add)집니다. 여기서는 step()과 zero_grad()을 호출하지 않앗기때문에 backward을 호출할때마다 gradient값이 누적됩니다.
• (step+1) % accumlation_steps == 0: 조건식을 만족하면 누적되어 평균된 gradient값에 대해서 모델의 가중치를 업데이트(step)하고, 이후 step에 반영되지 않도록 0으로 초기화해줍니다(zero_grad())

accumulation_steps = 4

model.train()
optimizer.zero_grad()

for step, (input_data, target_data) in enumerate(dataloader):
    outputs = model(input_data)
    loss = criterion(outputs, target_data)
    loss = loss / accumulation_steps
    loss.backward()

    if (step + 1) % accumulation_steps == 0:
        optimizer.step()
        optimizer.zero_grad()

 

수치적인 예시

---

예를 들어, 큰 배치(8)을 작은 배치(n=8)로 나누는 예시를 생각해볼 수 있습니다. 이 때, accumulation_steps은 2가 됩니다.

첫 번재 미니배치에서는 4개의 샘플이 들어옵니다

$loss_{1} = \frac{L_{1} + L_{2} + L_{3} + L_{4}}{4}$: 첫 번째 미니배치의 손실값입니다. pytorch에서 reduce을 미리 해주기에 평균값으로 자동으로 계산됩니다. 하지만, 최종적으로 8개 샘플에 대한 평균을 맞춰야 하므로, accumulation_steps = 2로 설정했다면, 아래와 같습니다.

$loss'_{1}  = \frac{1}{2} \times \frac{L_{1} + L_{2} + L_{3} + L_{4}}{4}$ 

이 때, loss_prime_1.backward()을 수행하면 $\nabla (loss_{1})/2$가 됩니다.

두 번째, 미니배치에서도 4개의 샘플이 들어옵니다.

$loss_{2} = \frac{L_{5} + L_{6} + L_{7} + L_{8}}{4}$: 두 번째 미니배치의 손실값입니다. 여기서도 accumulation_steps = 2로 설정햇기에 아래와 같습니다.

$loss'_{2}  = \frac{1}{2} \times \frac{L_{5} + L_{6} + L_{7} + L_{8}}{4}$ 

여기서도 loss_prime_2.backward()을 수행하면 $\nabla (loss_{2})/2$ 가 더해집니다.

 

최종적인 손실함수은 아래와 같습니다.

$\frac{1}{2} \times \frac{L_{1} + L_{2} + L_{3} + L_{4}}{4} + \frac{1}{2} \times \frac{L_{5} + L_{6} + L_{7} + L_{8}}{4} = \frac{L_{1} + ... + L_{8}}{8}$ 

 

그리고 최종적인 gradient은 아래와 같습니다. 즉, loss을 계산한후에 accumulation_step으로 나눠주어야, 원하는 큰 배치사이즈에 맞춰서 같은 스케일로 gradient을 업데이트할 수 있습니다.

$\frac{1}{2} \times \frac{\nabla (L_{1} + L_{2} + L_{3} + L_{4})}{4} + \frac{1}{2} \times \frac{\nabla (L_{5} + L_{6} + L_{7} + L_{8})}{4}  = \frac{ \nabla  (L_{1} + ... + L_{8})}{8} $