PyTorch에서 "CUDA out of memory" 에러를 방지하는 방법

2024-07-27

PyTorch에서 "CUDA out of memory" 에러를 방지하는 방법

데이터 크기 조절

배치 크기 줄이기: 배치 크기는 한 번에 학습에 사용되는 데이터의 양입니다. 배치 크기를 줄이면 GPU 메모리에 사용되는 데이터 양이 감소하여 에러 발생 가능성을 낮출 수 있습니다.
데이터 나누기: 학습 데이터가 매우 크다면 여러 개의 파트로 나누어 학습을 진행할 수 있습니다.
데이터 샘플링: 데이터 샘플링을 통해 학습에 사용되는 데이터의 양을 줄일 수 있습니다.
데이터 전처리: 이미지 크기 축소, 데이터 형식 변환 등의 전처리 과정을 통해 데이터 메모리 사용량을 줄일 수 있습니다.

모델 크기 조절

모델 파라미터 수 줄이기: 모델의 파라미터 수가 많을수록 더 많은 메모리가 필요합니다. 모델 구조를 변경하거나 정규화 기법을 사용하여 파라미터 수를 줄일 수 있습니다.
모델 병렬화: 모델 병렬화 기술을 사용하면 여러 GPU에 모델을 분산하여 학습시킬 수 있습니다. 이는 모델 크기에 따른 메모리 부족 문제를 해결하는 데 도움이 될 수 있습니다.

그래디언트 데이터 크기 조절

Grad scaler 사용: Grad scaler는 학습 과정에서 발생하는 그래디언트의 크기를 조절하는 도구입니다. Grad scaler를 사용하여 그래디언트 크기를 줄이면 메모리 사용량을 줄일 수 있습니다.
Mixed precision training: Mixed precision training은 학습 과정에서 float32와 float16 데이터 형식을 혼합하여 사용하는 기법입니다. 이 기법을 사용하면 메모리 사용량을 줄일 수 있지만, 모델 정확도가 약간 저하될 수 있습니다.

GPU 메모리 관리

torch.cuda.empty_cache() 사용: 학습 과정을 시작하기 전에 torch.cuda.empty_cache() 함수를 사용하여 GPU 메모리를 초기화하면 메모리 누수 문제를 방지할 수 있습니다.
CUDA 메모리 할당량 설정: torch.cuda.set_per_process_memory_fraction() 함수를 사용하여 프로세스별 GPU 메모리 할당량을 설정할 수 있습니다.

CPU 사용

GPU 메모리 부족 문제가 심각한 경우, CPU를 사용하여 학습을 진행할 수도 있습니다. CPU 학습 속도는 GPU보다 느리지만, 메모리 부족 문제를 해결할 수 있습니다.

기타

코드 최적화: 코드를 최적화하여 메모리 사용량을 줄일 수 있습니다.
데이터 증강: 데이터 증강 기법을 사용하여 학습 데이터를 늘릴 수 있습니다. 이는 모델 정확도를 향상시키는 동시에 메모리 사용량을 줄이는 데 도움이 될 수 있습니다.

예제 코드

import torch

# 데이터 크기 조절
batch_size = 16
data = torch.randn(batch_size, 100, 100)

# 모델 크기 조절
model = torch.nn.Sequential(
    torch.nn.Linear(100, 10),
    torch.nn.ReLU(),
    torch.nn.Linear(10, 1)
)

# 그래디언트 데이터 크기 조절
optimizer = torch.optim.SGD(model.parameters(), lr=0.01)
grad_scaler = torch.cuda.amp.GradScaler()

# GPU 메모리 관리
torch.cuda.empty_cache()

# 학습
for epoch in range(10):
    outputs = model(data)
    loss = torch.nn.MSELoss()(outputs, target)

    # Grad scaler 사용
    grad_scaler.scale(loss)
    optimizer.zero_grad()
    loss.backward()
    grad_scaler.unscale_(optimizer)
    optimizer.step()

    # CUDA 메모리 확인
    print(torch.cuda.memory_allocated())