tf.data 内存泄漏

Question

我在tf.data.Datasetfor 循环内部创建了一个，我注意到在每次迭代后内存并没有像人们期望的那样被释放。

有没有办法从 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

数据集 API 通过内置迭代器处理迭代，至少在 Eager 模式关闭或 TF 版本不是 2.0 时是这样。因此，根本不需要在 for 循环内从 numpy 数组创建数据集对象，因为它将图中的值写入为tf.constant. data = 的情况并非如此tf.data.TFRecordDataset()，因此如果您将数据转换为 tfrecords 格式并在 for 循环内运行它，则不会泄漏内存。

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