Tensorflow 사용시 메모리 누수가 발생하는 경우가 있어 오랬동안 고생하다 최근에 찾은 솔루션을 정리한다. 방법은 시스템 malloc 을 TCMalloc 으로 교체하는 것이다.

일부 시스템에서의 malloc 이 Tensorflow 와 함께 사용될 경우 반복적인 시스템 메모리에 접근하는 상황에서 메모리 회수가 잘 안되는 것 같다. 내 경우엔 다른 머신간에 다른 양상이 나타났는데, ① 메모리 누수가 안 발생, ② 시스템 메모리가 꽉 차서 해당 프로세스 kill, ③ 시스템 메모리가 꽉 차서 PC 먹통 등으로 나타났다(3번 경우 때문에 재택근무 중 사무실 출근이 필요했음).

시스템 malloc 을 TCMalloc 으로 간단하게 대체하기 위해서 LD_PRELOAD 트릭을 사용한다. LD_PRELOAD 는 환경변수로 넣어주는데, 프로그램이 로드되기 직전에 LD_PRELOAD 로 지정한 공유 라이브러리를 먼저 로드한다. 이렇게 함으로써 해당 라이브러리 심볼이 우선권을 갖게 된다. TCMalloc 이 정의된 공유라이브러리를 지정하면 실행하는 프로그램에서 malloc 이 호출되는 부분이 TCMalloc 으로 대체되는(override) 것이다.

우선 타겟 환경(도커 등)에 TCMalloc 을 설치한다.

$ apt install libtcmalloc-minimal4

그리고 타겟 환경에서 아래 명령어를 사용하여 libtcmalloc-minimal4.so.4 의 경로를 확인한다.

$ dpkg -L libtcmalloc-minimal4

그리고 타겟 환경에서 명령어를 실행할 때 LD_PRELOAD 환경변수를 지정해주면 된다.

$ LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4 train.py

내 경우에 메모리 누수현상은 사라졌지만 실행속도에서 큰 개선을 확인할 수는 없었다.

이 방법은 시스템 메모리 접근에 관련한 것이기 때문에 코드 자체를 tf.Function 으로 만 사용할 경우 누수 이슈가 안 나타날 수 있고 성능개선도 기대할 수 있을 것 같다. 내 경우에는 데이터셋 루프에서 시스템(cpu) 메모리로 가져와서 처리해야 편한 부분이 있어 이 방법이 현실적이고 빠른 해결책이다.

참고

• https://dhdroid.github.io/tensorflow/2021/10/11/detecting-tensorflow-memory-leak.html
• https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033
• https://stackoverflow.com/questions/426230/what-is-the-ld-preload-trick

매 번 다시 검색해서 알아내는 것이라 여기에 정리해둔다.

Tensorflow (버전 2.2와 1.14에서 확인)

import tensorflow as tf
tf.test.is_gpu_available(
    cuda_only=False, min_cuda_compute_capability=None
)

PyTorch (버전 1.6에서 확인)

import torch
torch.cuda.is_available()

참고:

• tf.test.is_gpu_available
• PyTorch GPU check