CODEWITHSUNDEEP
pythontensorflowmemory-leakstensorflow-datasets
BiBi
BiBi

Reputation: 7908

Memory leak with tf.data

I'm creating a tf.data.Dataset inside a for loop and I noticed that the memory was not freed as one would expect after each iteration.

Is there a way to request from TensorFlow to free the memory?

I tried using tf.reset_default_graph(), I tried calling del on the relevant python objects but this does not work.

The only thing that seems to work is gc.collect(). Unfortunately, gc.collect does not work on some more complex examples.

Fully reproducible code:

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)')

Evolution of memory usage

Upvotes: 9

Views: 7867

Answers (5)

Tonnz
Tonnz

Reputation: 43

If you only need to create then save the dataset to disk in the loop body, such as when wanting to preprocess large amounts of data in smaller parts to avoid out-of-memory, launch the loop body in a subprocess.

This answer describes how to launch subprocesses in general.

Upvotes: 0

joakimedin
joakimedin

Reputation: 21

This fix worked for me when I had a similar issue with TF 2.4

  1. sudo apt-get install libtcmalloc-minimal4
  2. LD_PRELOAD=/path/to/libtcmalloc_minimal.so.4 python example.py

Upvotes: 2

Dpk
Dpk

Reputation: 31

The issue is that you're adding a new node to the graph to define the iterator after each iteration, a simple rule of thumb is never define new tensorflow variables inside a loop. To fix it move 

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()

outside the for loop and just call sess.run(next_element) to fetch the next example and once youve gone through all the training/eval examples call sess.run(data_it) to reinitialize the iterator.

Upvotes: 3

MPękalski
MPękalski

Reputation: 7103

You are creating new python object (dataset) ever iteration of a loop and looks like garbage collector is not being invoked. Add impplicit garbage collection call and the memory usage should be fine.

Other than that, as mentioned in other answer, keep building data obect and session outside of the loop.

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)')

enter image description here

Upvotes: -1

Sharky
Sharky

Reputation: 4543

Dataset API handles iteration via built-in iterator, at least while eager mode is off or TF version is not 2.0. So, there's simply no need to create dataset object from numpy array inside for loop, as it writes values in the graph as tf.constant. This is not the case with data = tf.data.TFRecordDataset(), so if you transform your data to tfrecords format and run it inside for loop it won't leak memory.

for i in range(500):
    data = tf.data.TFRecordDataset('file.tfrecords')\
        .prefetch(64)\
        .repeat(-1)\
        .batch(1)
    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()

But as I said, there's no need to create dataset inside a loop.

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:
        ...

Upvotes: -1

Related Questions