純粹搬運工,接受英語的請看原網址:Keras Tensorflow backend automatically allocates all GPU memory。
通過設定Keras的Tensorflow後端的全局變量達到。
import os
import tensorflow as tf
import keras.backend.tensorflow_backend as KTF
def get_session(gpu_fraction=):
'''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 KTF
KTF.set_session(get_session())