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: