PyTorch GPU 메모리 해제하기 (PyTorch OOMError)

Pytorch로 모델을 학습하다보면 GPU 메모리가 꽉 차서 inference 등의 다른 작업이 안되는 경우가 있다. 기본적으로 메모리에 아무것도 올라가지 않은 상태에서 모델을 메모리에 올려 학습하려 할 때 나는 OOM은 배치 사이즈를 작게 조절하거나, 모델의 크기를 줄이는 등의 근본적인 문제를 해결해야 한다. 허나, 동일 조건에서 학습은 충분히 잘 되는데도 불구하고, 다른 작은 작업이 안되는 경우, 다음과 같은 방법을 통해 문제를 해결해볼 수 있다.

 

https://github.com/pytorch/pytorch/issues/16417 를 참고하였다.

 

Python Garbage Collection 기능 및 토치 캐시 비우기 기능을 사용한다. (참고 위치) 나의 경우, 아래 코드만 가지고 추론에 성공하였는데, 이것도 안되는 경우 위 링크의 다른 조언을 참고하거나, 위에서 언급한 것처럼 가지고 있는 리소스에서 원하는 모델을 학습할 수 없을 수 있으니 이 부분도 확인해보자.

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