How to control memory while using Keras with tensorflow backend?

Question

I have created a wrapper class which initializes a keras.models.Sequential model and has a couple of methods for starting the training process and monitoring the progress. I instantiate this class in my main file and perform the training process. Fairly mundane stuff.

My question is:

How to free all the GPU memory allocated by tensorflow. I tried the following with no luck:

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)

Even after the session has been closed and reset to None, nvidia-smi does not reflect any reduction in memory usage. Any ideas?

Idea

Would it be meaningful to add a __exit__ method to my NeuralNet class and instantiate it as:

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

How should I free up the resources of the keras model in this method?

Test environment

I'm using iPython Notebook on an Ubuntu 14.04 with 3 GTX 960 GPUs.

Reference:

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

Answer 1

The following works for me to reinitialize the state of Keras layers in my Jupyter notebook for every run:

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

Also, the graph is named and reset every time the notebook runs using:

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

Note : I am still trying to wrap my head around the concepts of Keras and tensorflow ( I believe they are described poorly in documentation and examples ) but the above works.