如何在使用带有 tensorflow 后端的 Keras 时控制内存？

Question

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

我的问题是：

如何释放tensorflow分配的所有GPU内存。我尝试了以下方法，但没有成功：

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 和 tensorflow 的概念（我相信它们在文档和示例中的描述不佳），但上述方法有效。