Keras设定GPU使用内存大小(Tensorflow backend)

通过设置Keras的Tensorflow后端的全局变量达到。

import osimport tensorflow as tfimport keras.backend.tensorflow_backend as KTFdef 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))

使用过程中显示的设置session:

import keras.backend.tensorflow_backend as KTFKTF.set_session(get_session())