[keras] GPU 메모리 사용량 제한하기

https://groups.google.com/forum/#!topic/keras-users/MFUEY9P1sc8

위 링크에서 찾을 수 있었는데

정리하자면 아래의 소스를 처음 부분에 입력해주면 GPU 메모리 전체를 사용하지 않는다. gpu_fraction 의 값에 따라서 전체 메모리 중 몇 퍼센트를 이용할 것인지 설정할 수 있다. 

Python 3.5, Keras 2.1.2 버전에서 실행하였다. 

import os
import tensorflow as tf
import keras.backend.tensorflow_backend as KK


def get_session(gpu_fraction=0.3):
    '''Assume that you have 6GB of GPU memory and want to allocate ~2GB'''

    num_threads = os.environ.get('OMP_NUM_THREADS')
    gpu_options = tf.GPUOptions(per_process_gpu_memory_fraction=gpu_fraction)

    if num_threads:
        return tf.Session(config=tf.ConfigProto(
            gpu_options=gpu_options, intra_op_parallelism_threads=num_threads))
    else:
        return tf.Session(config=tf.ConfigProto(gpu_options=gpu_options))


KK.set_session(get_session())