在使用TensorFlow作为backend的情况下控制Keras的显存使用(ConfigProto)

之前的一篇文章介绍了如何控制TensorFlow的显存使用,避免默认将显存吃满,可以通过ConfigProto的设置来控制。 当使用Keras的情况下,当import keras时,框架会自动开启一个默认参数的Session。可以通过下面的代码主动创建一个使用了ConfigProto的Session,再注入到Keras的后端中去: from keras.backend.tensorflow_backend import set_session config = tf.ConfigProto() config.gpu_options.allow_growth = True set_session(tf.Session(config=config)) 注意:

TensorFlow与显存

开始学习用TensorFlow,这东西与Theano不同,默认情况下, 在开启一个Session后,几乎占用显卡的所有显存。如果同一个机器、显卡多个人使用,基本上就是先到先得,后来的程序会崩溃。查了下文档有两种方法控制显存: 第一种是设置成预加载比例: tf_config = tensorflow.ConfigProto() tf_config.gpu_options.per_process_gpu_memory_fraction = 0.5 # 分配50% session = tensorflow.Session(config=tf_config)