Question

我正在for循环中创建一个tf.data.Dataset，我注意到每次迭代后内存释放都没有达到预期。

是否可以从TensorFlow请求释放内存？

我尝试使用tf.reset_default_graph()，尝试在相关的python对象上调用del，但这不起作用。

似乎唯一起作用的是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

当我在 TF 2.4 中遇到类似问题时

This fix 为我工作

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

Answer 2

您每次循环迭代时都在创建新的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 3

问题是您要在图上添加一个新节点以定义每次迭代后的迭代器，一个简单的经验法则是永远不要在循环内定义新的tensorflow变量。要解决此问题

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循环外，只需调用sess.run（next_element）即可获取下一个示例，并在完成所有训练/评估示例后，调用sess.run（data_it）重新初始化迭代器。

Answer 4

Dataset API通过内置的迭代器处理迭代，至少在急切模式关闭或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:
        ...

tf.data发生内存泄漏

4 个答案: