Memory Leak with Tensorflow Experimental Save

Viewed 87

I have a loop where I am creating tensorflow datasets and then saving to directories for later use.

I found that as the loop progresses, the memory being used greatly increases, until the process eventually crashes.

In this first example, I do not save the file, and the memory stays at the same level throughout:

import tensorflow as tf
import numpy as np
from humanize import naturalsize
import psutil

for i in range(10000):
    if not i % 100:
        print(naturalsize(psutil.Process().memory_info().rss))
    data = tf.data.Dataset.from_tensors(np.array([0, 1, 2]))

# prints 567.7 MB throughout

However, if I save each dataset, the memory used increases each time.

import tensorflow as tf
import numpy as np
from humanize import naturalsize
import psutil

for i in range(10000):
    if not i % 100:
        print(naturalsize(psutil.Process().memory_info().rss))
    data = tf.data.Dataset.from_tensors(np.array([0, 1, 2]))
    tf.data.experimental.save(data, path='~/Desktop/1')
# Grows by about 5 MB each 100 runs.

Adding tf.keras.backend.clear_session() after each save appears to slow down the memory growth but doesn't fully stop it.

Thank you in advance for any help.

0 Answers
Related