PyTorch CUDA out of memory 에러 완벽 해결 가이드 메모리 누수 원인과 캐시 정리 방법

PyTorch로 모델을 학습하다가 "RuntimeError: CUDA out of memory" 에러를 마주했다면, 단순히 GPU 성능 문제가 아닐 가능성이 높습니다. 실제로는 코드 레벨에서 발생하는 메모리 누수(Memory Leak)가 주범인 경우가 대부분입니다. 이 글에서는 NVIDIA GPU 환경(Jetson 포함)에서 발생하는 OOM 문제의 근본 원인을 진단하고, 즉시 적용 가능한 해결책을 제시합니다.


PyTorchCUDAoutofmemory에러완벽해결가이드메모리누수원인과캐시정리방법



1. CUDA OOM 에러의 두 가지 원인

GPU 메모리 부족 문제는 크게 두 가지 원인으로 나뉩니다.

  • 하드웨어 한계: 모델 크기 또는 배치(Batch) 사이즈가 GPU 메모리 용량을 초과하는 경우
  • 코드 레벨 메모리 누수: Autograd 그래프 참조 유지, 텐서 삭제 미흡, 캐시 미해제 등으로 인한 메모리 반환 실패

실무에서는 후자가 더 빈번합니다. 특히 반복문(iteration) 내에서 numpy 배열을 텐서로 변환하거나, 모델 출력값을 다른 함수에 직접 전달할 때 문제가 발생합니다.

2. 메모리 누수가 발생하는 전형적인 패턴

패턴 1: 반복문에서 텐서 생성 후 미흡한 정리

numpy 배열을 반복적으로 텐서로 변환할 때, del로 변수를 삭제해도 GPU 메모리는 즉시 반환되지 않습니다. PyTorch는 성능 최적화를 위해 캐시를 유지하기 때문입니다.

  • 문제 코드:
for epoch in range(num_epochs):
    trainData = torch.tensor(numpyTrainData).to(device)
    # ... 학습 코드 ...
    del trainData  # GPU 캐시는 여전히 남아있음
  • 해결 코드:
for epoch in range(num_epochs):
    trainData = torch.tensor(numpyTrainData).to(device)
    # ... 학습 코드 ...
    del trainData
    torch.cuda.empty_cache()  # 캐시 명시적 해제

패턴 2: 모델 출력값의 Autograd 그래프 참조 유지

모델의 출력(output)을 다른 함수나 변수에 직접 할당하면, Autograd가 계산 그래프를 유지하여 메모리가 해제되지 않습니다.

  • 문제 코드:
output = model(input_data)
result = some_function(output)  # 그래프 참조 유지
  • 해결 코드:
output = model(input_data)
result = some_function(output.data.cpu().numpy())  # 명시적 복사
# 또는
result = some_function(output.detach().cpu())

3. 체계적인 문제 진단 및 해결 프로세스

Step 1: 현재 GPU 메모리 상태 확인

터미널에서 nvidia-smi 명령어를 실행하여 실시간 메모리 할당 현황을 모니터링합니다.

nvidia-smi
# 또는 실시간 모니터링
watch -n 1 nvidia-smi

출력 결과에서 "Memory-Usage" 항목을 확인하여 allocated/cached 메모리를 구분합니다.

Step 2: 메모리 누수 방지를 위한 코드 체크리스트

  1. 모델 출력 처리: 추론 결과를 다른 곳에서 사용할 때는 반드시 .detach() 또는 .data.cpu().numpy()로 그래프를 분리합니다.
  2. 불필요한 텐서 제거: 반복문 내에서 생성한 임시 텐서는 deltorch.cuda.empty_cache()를 호출합니다.
  3. 디바이스 설정 검증: GPU 사용 여부를 명시적으로 확인합니다.
device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
model.to(device)

Step 3: TensorFlow 환경에서의 메모리 제한 설정

TensorFlow를 사용하는 경우, 프로그램 시작 시점에 GPU 메모리 사용량을 제한할 수 있습니다. 이 설정은 반드시 런타임 초기화 전에 실행해야 합니다.

import tensorflow as tf

gpus = tf.config.experimental.list_physical_devices('GPU')
if gpus:
    try:
        tf.config.experimental.set_virtual_device_configuration(
            gpus[0],
            [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]
        )
    except RuntimeError as e:
        print(e)

4. 주의해야 할 함정과 실전 팁

  • torch.cuda.empty_cache()의 한계: 이 함수는 캐시된 메모리를 해제할 뿐, 실제로 참조가 남아있는 텐서는 해제하지 못합니다. 반드시 del이나 .detach()를 먼저 수행해야 합니다.
  • CUDA error: device-side assert triggered 오류: 이 에러는 GPU 문제처럼 보이지만, 실제로는 텐서 차원 불일치(dimension mismatch)가 원인인 경우가 많습니다. CrossEntropyLoss 사용 시 출력 차원(output.shape)과 클래스 개수, 타깃 인덱스 범위를 반드시 확인하십시오.
  • Jetson 환경 특이사항: Jetson Nano/Xavier 등 임베디드 GPU는 통합 메모리(Unified Memory)를 사용하므로, 시스템 메모리와 GPU 메모리를 함께 고려해야 합니다.
  • NVIDIA PhysX/FrameView SDK 충돌: 일부 환경에서는 NVIDIA 관련 백그라운드 프로세스가 GPU 메모리를 점유할 수 있습니다. 제어판에서 불필요한 NVIDIA 구성 요소를 제거하면 해결됩니다.

요약 및 Action Item

CUDA OOM 문제는 하드웨어 한계보다 코드 레벨의 메모리 관리 미흡에서 비롯되는 경우가 많습니다. 핵심은 (1) Autograd 그래프 참조를 명시적으로 끊고, (2) 불필요한 텐서를 삭제한 후 캐시를 비우며, (3) nvidia-smi로 실시간 메모리 상태를 모니터링하는 것입니다.

지금 당장 실행할 것: 학습 루프의 반복문 끝에 다음 두 줄을 추가하십시오.

del unnecessary_tensors
torch.cuda.empty_cache()

이 조치만으로도 대부분의 메모리 누수 문제가 해결됩니다.





# 함께 보면 좋은 글

댓글