GPU memory leakage when I load Keras models multi times

user_9229923 0
Viewed 297

I have a file train.py to train neural net model via tensorflow.keras, and the best model will be saved as best_train.h5 according to the train_loss. The training time is about 7 days.

I have another file test.py to test the model saved during the running of train.py. In the test.py model, I load the best_train.h5 every hour to see the test performance. The codes are as follows:

for i in range(7*24):
  time.sleep(1*60*60)
  model = tf.keras.models.load_model('best_train.h5')
  model.predict(test_data)

I found that every time I load best-train.h5, the occupied GPU memory will be increased. And after about 200 iterations, the GPU memory will be run out. OOM Errors will occur.

I know tf.clear_session() can free the GPU memory. But this command will clear all the sessions and graph. I is not what I wanted. Because in test.py there are still other models I hold.

keras tensorflow python
1 Answers

Since Keras shares the global session between all models. You can create a new graph and assign a session for prediction only:

self.graph = tf.Graph()
with self.graph.as_default():
   self.session = tf.Session(graph=self.graph)
   with self.session.as_default(): 
       # Load your model and preform prediction

Your GPU memory should be released as soon as you finish your prediction.

user_61128300
Related