如何在使用具有张量流后端的Keras时控制内存？

Question

我创建了一个初始化keras.models.Sequential模型的包装类,它有几种启动培训过程和监控进度的方法.我在我的main文件中实例化这个类并执行培训过程.相当普通的东西.

我的问题是:

如何释放所有分配的GPU内存tensorflow.我试了以下没有运气:

import keras.backend.tensorflow_backend as K
with K.get_session() as sess:
    K.set_session(sess)
    import tensorflow as tf
    from neural_net import NeuralNet
    with tf.device('/gpu:0'):
        nn = NeuralNet('config', train_db_path, test_db_path)
        nn.train(1000, 1)
        print 'Done'
    K._SESSION.close()
    K.set_session(None)

即使会话已关闭并重置为None,nvidia-smi也不会反映内存使用量的任何减少.有任何想法吗？

理念

将__exit__方法添加到我的NeuralNet类并将其实例化为:

with NeuralNet() as nn:
    nn.train(1000, 1)

我应该如何在这种方法中释放keras模型的资源？

测试环境

我在带有3个GTX 960 GPU的Ubuntu 14.04上使用iPython Notebook.

参考:

1. https://github.com/fchollet/keras/issues/2102
2. https://groups.google.com/forum/#!topic/keras-users/MFUEY9P1sc8

Answer 1

以下内容可帮助我在每次运行时重新初始化 Jupyter 笔记本中 Keras 层的状态：

from keras import backend as K
K.clear_session()
sess = tf.Session()
K.set_session(sess)

此外，每次笔记本运行时，都会使用以下命令命名并重置图表：

graphr = K.get_session().graph
with graphr.as_default():
    #...graph building statements...

注意：我仍在尝试理解 Keras 和张量流的概念（我相信它们在文档和示例中描述得不好），但上面的方法是有效的。