Утечка памяти с tf.data

Question

Я создаю tf.data.Dataset внутри цикла for и заметил, что память не освобождалась, как можно было ожидать после каждой итерации.

Есть ли способ запросить у TensorFlow освободить память?

Я пытался использовать tf.reset_default_graph(), я пытался вызвать del для соответствующих объектов Python, но это не работает.

Единственное, что, кажется, работает - gc.collect().К сожалению, gc.collect не работает на некоторых более сложных примерах.

Полностью воспроизводимый код:

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

Answer 1

Вы создаете новый объект Python (набор данных) при каждой итерации цикла и похоже, что сборщик мусора не вызывается.Добавьте неявный вызов сборки мусора, и использование памяти должно быть в порядке.

Кроме этого, как упоминалось в другом ответе, оставляйте данные об объектах и ​​сеанс вне цикла.

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

Answer 2

API набора данных обрабатывает итерацию с помощью встроенного итератора, по крайней мере, когда режим ожидания выключен или версия TF не 2.0. Таким образом, просто не нужно создавать объект набора данных из массива numpy внутри цикла for, поскольку он записывает значения в графе как tf.constant. Это не относится к данным data = tf.data.TFRecordDataset(), поэтому, если вы преобразуете свои данные в формат tfrecords и запускаете их для цикла, это не приведет к утечке памяти.

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

Но, как я уже сказал, нет необходимости создавать набор данных внутри цикла.

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