Tensorflow에서 GPU 사용량 조절하기

GPU가 장착된 피씨 환경에서 Tensorflow를 사용하여 네트워크를 디자인하고 학습용 Dataset을 이용해서 네트워크 학습을 진행하다 보면 종종 Out of Memory (OOM) 에러가 자주 발견된다.

학습 Dataset에 사용된 이미지의 해상도가 크거나 네트워크의 깊이가 깊고 필터의 개수가 많은 경우 등 여러가지의 원인이 있겠지만, 어쨌든 필요로 하는 모든 자원들이 GPU 메모리 내부에 저장되어야 하는데 더 이상 할당이 안될 경우 높은 확률로 OOM 에러가 발생되며 실행이 중단된다.

이 때, GPU 사용량을 조절하는 방법에는 크게 2가지가 있는데,

1. 처음 시작 시 GPU의 사용이 최소가 되도록 설정하고, 필요한 경우가 발생될 때 마다 조금씩 더 할당하도록 설정하는 방법

2. GPU 메모리의 사용 가능량의 최대치를 설정하고 그 이상은 사용을 안하도록 제한하는 방법. 이 경우 학습에 속도는 저하될 수 있으나 GPU 메모리 이외의 다른 리소스를 활용하기 때문에 학습이 느리더라도 어쨌든 진행은 되는 것을 확인 했다.

첫번째 방법은 allow_growth 옵션을 사용함으로써 적용이 가능하며,
아래와 같은 예제로 적용하면 된다.


config = tf.ConfigProto()


config.gpu_options.allow_growth = True


session = tf.Session(config=config,...)



두번째 방법은 per_process_gpu_memory_fraction 옵션을 통해 선언이 가능하며,
아래와 같은 방법으로 적용하면 된다.

config = tf.ConfigProto()
 
config.gpu_options.per_process_gpu_memory_fraction = 0.6 
 #GPU로 60%를 사용하고 싶을 때

session = tf.Session(config=config,...)