在使用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))  

注意:第一行必须直接import这个函数,而非keras这个包,而且这句话必须在其他import keras之前,否则keras初始化之后,再替换一个session,原session也不会释放。

Friskit

继续阅读此作者的更多文章