그래디언트 누적

메모리에 들어가지 않는 큰 배치

Transformer를 비롯한 일부 모델은 큰 배치에서 더 안정적으로 학습되지만, 원하는 배치 전체가 GPU 메모리에 들어가지 않을 수 있어. 그래디언트 누적은 통계적인 배치와 물리적인 배치를 분리하는 방법이야. 작은 마이크로배치를 여러 번 처리하면서 그래디언트를 모으고, 모인 값이 목표 배치를 대표할 때 옵티마이저를 한 번 갱신해. 그러면 가중치 갱신은 큰 배치를 한 번 처리한 결과와 거의 같아져.

계산의 의미를 보존하는 두 조건

첫째, 각 마이크로배치의 손실을 accumulation_steps로 나눠야 누적 그래디언트가 합계가 아니라 평균이 돼. 나누지 않으면 유효 학습률이 누적 횟수만큼 커져 버려. 둘째, 적용과 초기화를 같은 경계에서 수행해 옵티마이저 한 단계가 정확히 N개의 마이크로배치를 소비하도록 해야 해. 아래 코드는 TensorFlow 그래디언트 기록 장치로 만든 수동 루프야. 누적기를 모델 속성으로 옮기고 train_step()에 담으면 운영 코드에서도 사용할 수 있어.

Code

그래디언트 누적(TensorFlow 백엔드)·python

import tensorflow as tf

accumulation_steps = 4  # Effective batch = batch_size x 4
accumulated = None

for step, (x, y) in enumerate(dataset):
    with tf.GradientTape() as tape:
        y_pred = model(x, training=True)
        # Divide so the summed gradient is an AVERAGE
        loss = loss_fn(y, y_pred) / accumulation_steps

    grads = tape.gradient(loss, model.trainable_variables)

    # Accumulate
    if accumulated is None:
        accumulated = grads
    else:
        accumulated = [a + g for a, g in zip(accumulated, grads)]

    # Apply every N micro-batches, then reset
    if (step + 1) % accumulation_steps == 0:
        optimizer.apply_gradients(
            zip(accumulated, model.trainable_variables)
        )
        accumulated = None

메모리에 들어가지 않는 큰 배치

계산의 의미를 보존하는 두 조건

Code

External links

Exercise

Progress

댓글 0