Gradient Accumulation

RuntimeError: CUDA out of memory.

메모리 용량이 부족해 발생하는 문제이기 때문에 단순히 batch size를 줄이는 등 공간을 확보해 해결할 수 있다. 또 의도적으로 GC를 실행해 메모리를 확보하거나 cache를 비워줄 수 있다. 

import gc
import torch

gc.collect()
torch.cuda.empty_cache()

하지만 batch 크기를 줄이게 되면 성능이 떨어지는 문제가 발생할 수 있다. 따라서 큰 batch 크기를 사용해 학습하는 것과 동일한 효과를 내는 방법이 존재한다. 

---

Gradient Accumulation

기존의 학습은 하나의 mini-batch를 거친 후 optimizer를 통해 학습을 진행한다. 하지만 gradient accumulation는 여러 개의 mini-batch를 거친 후 누적된 값을 한 번에 학습한다. 

이렇게 하면 batch 크기를 키운 것과 유사한 효과를 얻을 수 있다. 

 

기존의 방식:

for epoch in range(1, 5):
    model.train()
    for data in train_set:
        inputs, label = data
        output = model(**inputs)    

        batch_loss = criterion(output.logits, label) 
        batch_loss.backward()

        optimizer.step()
        optimizer.zero_grad()

하나의 배치 당 한 번의 step이 진행된다. 

 

Gradient Accumulation:

for epoch in range(1, 5):
    # batch_size = 8
    # grad_step = 8

    model.train()
    for batch_id, data in enumerate(train_set, start=1):
        inputs, label = data
        output = model(**inputs)    

        batch_loss = criterion(output.logits, label) 
        (batch_loss / grad_step).backward()

        if batch_id % grad_step == 0:
            optimizer.step()
            optimizer.zero_grad()

grad_step번의 batch를 거친 후에  한 번의 step이 진행된다.