7
votes

Fuite de mémoire avec tf.data

Je crée un tf.data.Dataset dans une boucle for et j'ai remarqué que la mémoire n'était pas libérée comme on pourrait s'y attendre après chaque itération.

Y a-t-il un moyen de demander de TensorFlow pour libérer la mémoire?

J'ai essayé d'utiliser tf.reset_default_graph () , j'ai essayé d'appeler del sur les objets python concernés mais ce n'est pas le cas fonctionne.

La seule chose qui semble fonctionner est gc.collect () . Malheureusement, gc.collect ne fonctionne pas sur certains exemples plus complexes.

Code entièrement reproductible:

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import psutil
%matplotlib inline

memory_used = []
for i in range(500):
    data = tf.data.Dataset.from_tensor_slices(
                    np.random.uniform(size=(10, 500, 500)))\
                    .prefetch(64)\
                    .repeat(-1)\
                    .batch(3)
    data_it = data.make_initializable_iterator()
    next_element = data_it.get_next()

    with tf.Session() as sess:
        sess.run(data_it.initializer)
        sess.run(next_element)
    memory_used.append(psutil.virtual_memory().used / 2 ** 30)
    tf.reset_default_graph()

plt.plot(memory_used)
plt.title('Evolution of memory')
plt.xlabel('iteration')
plt.ylabel('memory used (GB)')

python tensorflow tensorflow-datasets memory-leaks

0 commentaires

3 Réponses :

-1
votes

L'API du jeu de données gère l'itération via un itérateur intégré, au moins lorsque le mode impatient est désactivé ou que la version TF n'est pas 2.0. Ainsi, il n'est tout simplement pas nécessaire de créer un objet de jeu de données à partir du tableau numpy à l'intérieur de la boucle for, car il écrit les valeurs dans le graphique sous la forme tf.constant . Ce n'est pas le cas avec data = tf.data.TFRecordDataset () , donc si vous transformez vos données au format tfrecords et que vous les exécutez à l'intérieur de la boucle for, cela ne fuira pas de mémoire.

data = tf.data.Dataset.from_tensor_slices(
                    np.random.uniform(size=(10, 500, 500)))\
                    .prefetch(64)\
                    .repeat(-1)\
                    .batch(3)
data_it = data.make_initializable_iterator()
next_element = data_it.get_next()

for i in range(500):
    with tf.Session() as sess:
        ...

2 commentaires

Je connais bien l'API tf.data , ma question porte sur un autre point, à savoir libérer explicitement la mémoire allouée par TensorFlow avec tf.data .

Faites une boucle à ceci, cela montre spécifiquement pourquoi ce n'est pas une bonne idée d'essayer de libérer de la mémoire github .com / tensorflow / tensorflow / issues / 14181

0
votes

Vous créez un nouvel objet Python (ensemble de données) à chaque itération d'une boucle et il semble que le garbage collector n'est pas appelé. Ajoutez un appel de garbage collection impplicite et l'utilisation de la mémoire devrait être correcte.

À part cela, comme mentionné dans une autre réponse, continuez à créer un objet de données et une session en dehors de la boucle.

import tensorflow as tf
import numpy as np
import matplotlib.pyplot as plt
import psutil
import gc

%matplotlib inline

memory_used = []
for i in range(100):
    data = tf.data.Dataset.from_tensor_slices(
                    np.random.uniform(size=(10, 500, 500)))\
                    .prefetch(64)\
                    .repeat(-1)\
                    .batch(3)
    data_it = data.make_initializable_iterator()
    next_element = data_it.get_next()

    with tf.Session() as sess:
        sess.run(data_it.initializer)
        sess.run(next_element)
    memory_used.append(psutil.virtual_memory().used / 2 ** 30)
    tf.reset_default_graph()
    gc.collect()

plt.plot(memory_used)
plt.title('Evolution of memory')
plt.xlabel('iteration')
plt.ylabel('memory used (GB)')

entrez la description de l'image ici a >

5 commentaires

Désolé, je viens de remarquer que vous avez écrit que vous avez essayé gc.collect (). Mais quel serait le cas d'utilisation le plus complexe?

Oui, j'ai essayé gc.collect () . Mon cas d'utilisation plus complexe implique de nombreux fichiers * .tfrecord avec un pipeline de données assez complexe. Dans ce cas d'utilisation plus complexe, gc.collect () ne fonctionne pas. Ainsi ma question: comment libérer explicitement la mémoire allouée par TensorFlow.

J'ai lu quelque part à ce sujet et en général: avec une mémoire GPU, il n'est pas possible de la libérer, avec Python cela doit être gc.collect () mais vous avez de nombreux problèmes typiques avec Python qui le font ne veut pas libérer la mémoire - il y a une grande menace à ce sujet sur stackoverflow.

Ici, je fais référence à la RAM générale, pas à la mémoire du GPU.

Je pense que vous devez décrire votre cas d'utilisation plus précisément. Avez-vous besoin de parcourir tous ces différents fichiers TFRecord à la fois ou non?

0
votes

Le problème est que vous ajoutez un nouveau nœud au graphe pour définir l'itérateur après chaque itération, une règle empirique simple est de ne jamais définir de nouvelles variables tensorflow à l'intérieur d'une boucle. Pour résoudre ce problème, déplacez

data = tf.data.Dataset.from_tensor_slices(
            np.random.uniform(size=(10, 500, 500)))\
            .prefetch(64)\
            .repeat(-1)\
            .batch(3)
data_it = data.make_initializable_iterator()
next_element = data_it.get_next()

hors de la boucle for et appelez simplement sess.run (next_element) pour récupérer l'exemple suivant et une fois que vous avez parcouru tous les exemples d'entraînement / d'évaluation, appelez sess.run (data_it ) pour réinitialiser l'itérateur.

0 commentaires