[MLOps] PyTorch GPU Out-Of-Memory 문제 해결

PyTorch를 이용하여 ML Pipeline을 개발하던 중 특정 시간에 학습을 시킬 때 메모리가 꽉 차서 모델 학습할 때 실패하는 경우가 있었다. 

ML Pipeline을 구성할 때 메모리도 넉넉하게 사용할 수 있도록 구성하지만 의도치 않게 OOM문제가 발생하는 경우가 있었다. APScheduler를 이용해서 구성했는데 학습을 하면서 out-of-memory 이슈가 발생해서 문제를 해결하기 위해 찾아보았고, 같은 이유로 해결방법을 찾는 사람들에게 도움이 되고자 한다. 

 

파이토치 뿐만 아니라 텐서플로우를 사용하더라도 마찬가지이므로 도움이 될 것 이다. 

 

GPU Out-Of-Memory 발생

RuntimeError: CUDA out of memory. Tried to allocate 126.00 MiB (GPU 0; 16.2 GiB total capacith; 16.02 GiB already allocated; 34.25 MiB free; 13.25 GiB cached)

 

1회차 학습 때는 이런 문제가 발생하지 않았으나, 2회차 학습 때부터는 다음과 같은 이슈가 발생했다. 

배치나 학습에 관련된 configuration이 같을 때 이와 같이 발생하는 경우는 없었는데, APScheduler를 이용해서 지속적으로 학습을 시킬 경우에 발생하는 이유는  APScheduler에서 학습이 끝난 이후에 메모리 할당된 것을 free 시켜주지 않았기 때문이다.

 

APScheduler에서 학습을 진행한 후에 모든 메모리 할당을 취소해주어야한다. 

 

학습이 끝났을 때도 torch.cuda.empty_cache()로 메모리 할당을 취소해주지만, 

APScheduler를 이용해서 학습이 끝났을 때도 메모리 할당을 취소해주어야 한다. 

 

메모리 할당 제거

torch.cuda.empty_cache()

 

다음과 같이 torch cuda 메모리 할당을 제거해주면 정상적으로 원하는 시간에 학습이 끊기지 않고 잘 돌아갈 것이다.