딥러닝 모델을 실험하거나 배포할 때 가장 흔히 부딪히는 문제 중 하나가 VRAM 부족입니다. 고성능 GPU가 없더라도 보다 효율적으로 모델을 돌리기 위해, PyTorch 환경에서 VRAM을 절약할 수 있는 10가지 방법을 정리했습니다.
1. FP16(반정밀도) 사용
모델과 입력 데이터를 float16(반정밀도)로 변환하면 VRAM 사용량을 대폭 줄일 수 있습니다.
model.half()
data = data.half()하지만 연산 정확도 손실이 발생할 수 있으므로, 자동 혼합 정밀도(AMP)와 함께 사용하는 것이 좋습니다.
2. 배치 크기 줄이기
가장 간단하고 확실한 방법 중 하나입니다. 배치 크기를 줄이면 처리할 데이터가 줄어 VRAM 소모도 감소합니다.
batch_size = 1
test_loader = torch.utils.data.DataLoader(dataset, batch_size=batch_size, num_workers=0, shuffle=False)3. 모델 일부를 CPU로 이동
특정 모듈만 CPU에서 실행하도록 하면 GPU 메모리를 절약할 수 있습니다.
device_process = 'cpu'
x = Process(x).to(device_process)4. Gradient 계산 비활성화
학습이 아닌 추론 단계에서는 torch.no_grad()를 사용하여 불필요한 그래디언트 계산을 방지할 수 있습니다.
with torch.no_grad():
output = model(data)5. 불필요한 파라미터 제거
모델 구조를 분석하고 사용하지 않는 레이어나 파라미터를 제거하면 모델 크기와 메모리 사용량을 줄일 수 있습니다.
6. 입력 데이터 크기 줄이기
입력 이미지나 텐서의 해상도를 줄이면 연산량과 메모리 소모가 동시에 줄어듭니다.
data = torch.nn.functional.interpolate(data, scale_factor=0.5)7. Checkpoint 활용
PyTorch의 torch.utils.checkpoint를 활용하면 중간 결과를 저장하지 않고 필요할 때 다시 계산하여 메모리를 절약할 수 있습니다.
from torch.utils.checkpoint import checkpoint
output = checkpoint(model, data)8. 모델 압축: 양자화 & 프루닝
모델을 양자화(Quantization) 또는 프루닝(Pruning) 하면 VRAM 사용량은 물론 연산 속도도 개선됩니다.
9. CUDA 메모리 정리
훈련 루프 중간중간 다음 코드를 사용하면 사용되지 않는 메모리를 해제할 수 있습니다.
torch.cuda.empty_cache()10. 모델을 분리하여 실행
모델의 일부를 나누어 순차적으로 실행하는 방식도 VRAM 사용량을 조절하는 전략 중 하나입니다.
마무리
이처럼 여러 기법을 적절히 조합하면, VRAM 부족 문제를 효율적으로 해결할 수 있습니다. 특히 모델 디버깅이나 로컬 환경에서 실험할 때 위 방법들이 큰 도움이 될 것입니다.
