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.