python - que - ¿Cómo controlar la memoria mientras usas Keras con tensorflow back-end?

tensorflow python

Question

He creado una clase contenedora que inicializa un modelo keras.models.Sequential y tiene un par de métodos para iniciar el proceso de formación y supervisar el progreso. Instalé esta clase en mi archivo main y realizo el proceso de entrenamiento. Cosas bastante mundanas.

Mi pregunta es:

Cómo liberar toda la memoria de la GPU asignada por tensorflow . Intenté lo siguiente sin suerte:

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)

Incluso después de que la sesión se haya cerrado y restablecido a None , nvidia-smi no refleja ninguna reducción en el uso de la memoria. ¿Algunas ideas?

Idea

¿Sería significativo agregar un método __exit__ a mi clase NeuralNet e instanciarlo como sigue?

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

¿Cómo debería liberar los recursos del modelo keras en este método?

Entorno de prueba

Estoy usando iPython Notebook en Ubuntu 14.04 con 3 GPU GTX 960.

Referencia:

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

Answer 1

Lo siguiente me funciona para reiniciar el estado de las capas de Keras en mi portátil Jupyter para cada ejecución:

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

Además, el gráfico recibe su nombre y se restablece cada vez que se ejecuta utilizando:

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

Nota: todavía estoy tratando de entender los conceptos de Keras y tensorflow (creo que se describen mal en documentación y ejemplos), pero lo anterior funciona.