[정보] pytorch GPU out-of-memory 문제 해결하기

pytorch를 사용해서 모델을 개발하다 보면 모델에 문제는 없는데 자꾸 메모리가 꽉 차서 모델 학습이 실패하는 경우가 종종 있다. 이러한 경우 문제가 아주 사소한 부분에서 발생하기 때문에 해결 방법을 찾기가 무척이나 힘들다. 본 문서에서는 본인이 pytorch GPU out-of-memory 문제를 해결했던 방법을 정리하여 pytorch로 모델 개발 중 out-of-memory 문제로 고민하는 사람들에게 도움이 되고자 한다.

pytorch GPU out-of-memory문제가 발생하는 원인

Traceback (most recent call last):
  File "main.py", line 161, in <module>
    tran_eval.learn(mem)
  File "/home/admin/Desktop/Model/model.py", line 35, in learn
    loss.backward()
  File "/home/admin/anaconda3/envs/model/lib/python3.7/site-packages/torch/tensor.py", line 166, in backward
    torch.autograd.backward(self, gradient, retain_graph, create_graph)
  File "/home/admin/anaconda3/envs/model/lib/python3.7/site-packages/torch/autograd/__init__.py", line 99, in backward
    allow_unreachable=True)  # allow_unreachable flag
RuntimeError: CUDA out of memory. Tried to allocate 122.00 MiB (GPU 0; 3.95 GiB total capacity; 3.08 GiB already allocated; 44.44 MiB free; 280.89 MiB cached)

위와 같은 pytorch GPU out-of-memory 문제가 발생하는 원인은 여러 가지가 있다. 우선은 모델 자체의 크기가 너무 크거나, 동영상이나 3D point cloud를 학습하는 등 GPU에 비해 학습하려는 데이터의 크기가 큰 경우 out-of-memory 문제가 발생할 수 있다. 하지만 작은 데이터셋과 몇 층의 neural network만을 사용하였음에도 out-of-memory 문제가 발생하는 경우가 종종 있다. 이러한 경우는 작성한 pytorch 코드 상에서 메모리 누수가 발생하였을 가능성이 높다. 해당 경우 문제를 해결하는 방법을 살펴보도록 하자.

case 1. 모델의 output을 다른 함수/코드에서 사용하는 경우

보통 학습 iteration에서 한 iteration이 끝나면 이전 모델의 output은 자동으로 메모리 반환이 이루어지지만, 모델의 output을 다른 함수나 코드에서 사용하는 경우 메모리 반환이 이루어지지 않는다. 따라서 이러한 경우 output을 복사해서 다른 함수나 코드에서 사용하여야 한다.

predict = model(input_data)
copiedPredict = predict.data.cpu().numpy() # 결과를 복사
otherFunction(copiedPredict) # 다른함수에서 결과를 사용

case 2. 학습 데이터가 삭제되지 않는 경우

본인의 경우 위의 방법을 사용하였음에도 GPU 메모리 사용량이 줄어들지 않았다. 코드를 자세히 살펴보았더니 매 iteration마다 numpy array로부터 tensor를 생성하고 있었는데 이러한 tensor들이 문제일 것으로 생각하고 해당 tensor들의 메모리를 del 함수를 사용하여 반환하였음에도 문제가 해결되지 않았다. 찾아보았더니 del 함수를 사용하여도 GPU 메모리 상에서 cache로 계속 데이터를 가지고 있을 수 있다고 한다. 그래서 아래의 코드와 같이 torch.cuda.empty_cache 함수를 호출해야 GPU 메모리 상에 학습 데이터가 남는 문제를 해결할 수 있다.

trainData = torch.tensor(numpyTrainData)
...모델학습...
del trainData # 학습 데이터 삭제
torch.cuda.empty_cache() # GPU 캐시 데이터 삭제