这篇文章主要讲解了Keras如何设定GPU使用内存大小,内容清晰明了,对此有兴趣的小伙伴可以学习一下,相信大家阅读完之后会有帮助。
通过设置Keras的Tensorflow后端的全局变量达到。
import os import tensorflow as tf import keras.backend.tensorflow_backend as KTF 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))